Testing

Testing GraphQL Subscriptions: WebSockets, Async Generators, and Jest

HelpMeTest

HelpMeTest

— 8 min read
Testing GraphQL Subscriptions: WebSockets, Async Generators, and Jest

GraphQL subscriptions are the hardest part of a GraphQL API to test. Queries and mutations are request/response — you call them, they return a value, done. Subscriptions are a persistent, stateful, event-driven channel over WebSockets. The moment you introduce async iteration, connection lifecycle, and client reconnection logic, every testing pattern you know for REST or queries stops working.

This post goes through each layer: resolver unit tests using async generators, integration tests with a real WebSocket server and the graphql-ws client, mocking subscriptions in React component tests, and testing error paths and reconnection behavior. All examples use Jest (or Vitest — the APIs are identical here), Apollo Server v4, and graphql-ws.

Why Subscription Testing Is Hard

A GraphQL subscription resolver is an async generator function. It yields values over time in response to events — a pub/sub channel, a database change stream, a timer. Three things make this hard to test:

Async iteration. for await...of loops don't have a natural end unless you close the generator or throw into it. Tests hang if you forget to close the subscription.

WebSocket lifecycle. The transport layer (connect, subscribe, ping/pong, reconnect) is separate from the resolver logic. A bug in either layer looks the same from the outside.

Timing. Events arrive asynchronously. Test code that doesn't wait for the right number of events will either miss them or assert prematurely.

Each of these requires a different fix. Don't try to solve all three in the same test.

Unit Testing Subscription Resolvers (Async Generators)

The subscription resolver is just an async generator. Test it in isolation — no server, no WebSocket. This is the fastest feedback loop.

// src/resolvers/chat.ts
import { pubsub } from '../pubsub';

export const chatSubscription = {
  messageAdded: {
    subscribe: (_root: unknown, args: { roomId: string }) =>
      pubsub.asyncIterableIterator(`MESSAGE_ADDED:${args.roomId}`),
    resolve: (payload: { messageAdded: Message }) => payload.messageAdded,
  },
};

For unit tests, don't use your real pub/sub (Redis, etc.). Write a minimal AsyncIterator you control directly:

// test/helpers/mockIterator.ts
export function createMockIterator<T>() {
  const queue: T[] = [];
  const resolvers: Array<(value: IteratorResult<T>) => void> = [];
  let done = false;

  return {
    iterator: {
      [Symbol.asyncIterator]() { return this; },
      next(): Promise<IteratorResult<T>> {
        if (queue.length > 0) {
          return Promise.resolve({ value: queue.shift()!, done: false });
        }
        if (done) {
          return Promise.resolve({ value: undefined as any, done: true });
        }
        return new Promise(resolve => resolvers.push(resolve));
      },
      return(): Promise<IteratorResult<T>> {
        done = true;
        resolvers.forEach(r => r({ value: undefined as any, done: true }));
        return Promise.resolve({ value: undefined as any, done: true });
      },
    },
    push(value: T) {
      if (resolvers.length > 0) {
        resolvers.shift()!({ value, done: false });
      } else {
        queue.push(value);
      }
    },
    close() {
      done = true;
      resolvers.forEach(r => r({ value: undefined as any, done: true }));
    },
  };
}

Now write the unit test:

// test/resolvers/chat.test.ts
import { createMockIterator } from '../helpers/mockIterator';
import { chatSubscription } from '../../src/resolvers/chat';

jest.mock('../../src/pubsub', () => ({
  pubsub: {
    asyncIterableIterator: jest.fn(),
  },
}));

import { pubsub } from '../../src/pubsub';

describe('messageAdded subscription resolver', () => {
  it('yields resolved messages for a given room', async () => {
    const mock = createMockIterator<{ messageAdded: { id: string; text: string } }>();
    (pubsub.asyncIterableIterator as jest.Mock).mockReturnValue(mock.iterator);

    const iterator = chatSubscription.messageAdded.subscribe(null, { roomId: 'room-1' });

    // Push two events
    mock.push({ messageAdded: { id: '1', text: 'hello' } });
    mock.push({ messageAdded: { id: '2', text: 'world' } });
    mock.close();

    const results: unknown[] = [];
    for await (const payload of iterator) {
      const resolved = chatSubscription.messageAdded.resolve(payload as any);
      results.push(resolved);
    }

    expect(results).toEqual([
      { id: '1', text: 'hello' },
      { id: '2', text: 'world' },
    ]);
    expect(pubsub.asyncIterableIterator).toHaveBeenCalledWith('MESSAGE_ADDED:room-1');
  });

  it('filters messages by roomId — does not yield from wrong room', async () => {
    // This tests that subscribe is called with the correct topic key,
    // not that it filters at the resolver level (the pub/sub does that)
    const mock = createMockIterator();
    (pubsub.asyncIterableIterator as jest.Mock).mockReturnValue(mock.iterator);

    chatSubscription.messageAdded.subscribe(null, { roomId: 'room-99' });
    mock.close();

    expect(pubsub.asyncIterableIterator).toHaveBeenCalledWith('MESSAGE_ADDED:room-99');
  });
});

Key discipline: always call mock.close() or iterator.return() in every test path, including failure paths. Without it the for await loop never terminates and Jest reports a timeout.

Integration Testing with a Real Server and graphql-ws Client

Unit tests cover resolver logic. Integration tests cover the full pipeline: schema, server, transport. Use a real Apollo Server with an in-process WebSocket server and the graphql-ws client.

// test/integration/subscriptions.test.ts
import { createServer } from 'http';
import { ApolloServer } from '@apollo/server';
import { ApolloServerPluginDrainHttpServer } from '@apollo/server/plugin/drainHttpServer';
import { makeExecutableSchema } from '@graphql-tools/schema';
import { WebSocketServer } from 'ws';
import { useServer } from 'graphql-ws/lib/use/ws';
import { createClient } from 'graphql-ws';
import WebSocket from 'ws';
import { typeDefs } from '../../src/schema';
import { resolvers } from '../../src/resolvers';
import { pubsub } from '../../src/pubsub';

let httpServer: ReturnType<typeof createServer>;
let wsServer: WebSocketServer;
let apolloServer: ApolloServer;
let serverCleanup: { dispose: () => Promise<void> };
let port: number;

beforeAll(async () => {
  const schema = makeExecutableSchema({ typeDefs, resolvers });

  httpServer = createServer();
  wsServer = new WebSocketServer({ server: httpServer, path: '/graphql' });
  serverCleanup = useServer({ schema }, wsServer);

  apolloServer = new ApolloServer({
    schema,
    plugins: [
      ApolloServerPluginDrainHttpServer({ httpServer }),
      {
        async serverWillStart() {
          return {
            async drainServer() {
              await serverCleanup.dispose();
            },
          };
        },
      },
    ],
  });

  await apolloServer.start();

  await new Promise<void>(resolve => {
    httpServer.listen(0, () => resolve()); // port 0 = OS assigns free port
  });

  port = (httpServer.address() as any).port;
});

afterAll(async () => {
  await apolloServer.stop();
  wsServer.close();
  httpServer.close();
});

function makeClient() {
  return createClient({
    url: `ws://localhost:${port}/graphql`,
    webSocketImpl: WebSocket,
  });
}

it('receives published messages over WebSocket', async () => {
  const client = makeClient();
  const received: unknown[] = [];

  const subscription = client.iterate({
    query: `subscription { messageAdded(roomId: "room-1") { id text } }`,
  });

  // Start iterating in the background
  const iterating = (async () => {
    for await (const event of subscription) {
      received.push(event.data);
      if (received.length === 2) break; // exit after 2 messages
    }
  })();

  // Wait for the subscription to be established before publishing
  await new Promise(resolve => setTimeout(resolve, 50));

  pubsub.publish('MESSAGE_ADDED:room-1', { messageAdded: { id: '1', text: 'hello' } });
  pubsub.publish('MESSAGE_ADDED:room-1', { messageAdded: { id: '2', text: 'world' } });

  await iterating;
  await client.dispose();

  expect(received).toEqual([
    { messageAdded: { id: '1', text: 'hello' } },
    { messageAdded: { id: '2', text: 'world' } },
  ]);
});

The client.iterate() API returns an AsyncIterableIterator, making it straightforward to consume in tests. The key pattern is starting iteration in a background promise, then publishing events, then awaiting the iteration to complete.

Replace the setTimeout with a more robust wait if your pub/sub setup has latency — listen for the connected callback on the client:

await new Promise<void>(resolve => {
  client = createClient({
    url: `ws://localhost:${port}/graphql`,
    webSocketImpl: WebSocket,
    on: { connected: () => resolve() },
  });
});

Mocking Subscriptions in React Component Tests

When testing React components that use Apollo Client subscriptions, you don't want a real server. Use MockSubscriptionLink from @apollo/client/testing.

// src/components/ChatRoom.tsx
import { useSubscription, gql } from '@apollo/client';

const MESSAGE_ADDED = gql`
  subscription MessageAdded($roomId: String!) {
    messageAdded(roomId: $roomId) {
      id
      text
    }
  }
`;

export function ChatRoom({ roomId }: { roomId: string }) {
  const { data, loading, error } = useSubscription(MESSAGE_ADDED, {
    variables: { roomId },
  });

  if (loading) return <p>Connecting...</p>;
  if (error) return <p>Error: {error.message}</p>;
  return <p data-testid="message">{data?.messageAdded.text}</p>;
}
// test/components/ChatRoom.test.tsx
import { render, screen, waitFor } from '@testing-library/react';
import { ApolloClient, InMemoryCache, gql } from '@apollo/client';
import { ApolloProvider } from '@apollo/client';
import { MockSubscriptionLink } from '@apollo/client/testing';
import { ChatRoom } from '../../src/components/ChatRoom';

function renderWithClient(link: MockSubscriptionLink) {
  const client = new ApolloClient({ link, cache: new InMemoryCache() });
  return render(
    <ApolloProvider client={client}>
      <ChatRoom roomId="room-1" />
    </ApolloProvider>
  );
}

it('renders message when subscription event arrives', async () => {
  const link = new MockSubscriptionLink();
  renderWithClient(link);

  expect(screen.getByText('Connecting...')).toBeInTheDocument();

  link.simulateResult({
    result: { data: { messageAdded: { id: '1', text: 'hello from server' } } },
  });

  await waitFor(() => {
    expect(screen.getByTestId('message')).toHaveTextContent('hello from server');
  });
});

it('renders error state when subscription errors', async () => {
  const link = new MockSubscriptionLink();
  renderWithClient(link);

  link.simulateResult({
    error: new Error('Subscription failed'),
  });

  await waitFor(() => {
    expect(screen.getByText(/Subscription failed/)).toBeInTheDocument();
  });
});

it('updates when multiple events arrive in sequence', async () => {
  const link = new MockSubscriptionLink();
  renderWithClient(link);

  link.simulateResult({
    result: { data: { messageAdded: { id: '1', text: 'first' } } },
  });

  await waitFor(() => expect(screen.getByTestId('message')).toHaveTextContent('first'));

  link.simulateResult({
    result: { data: { messageAdded: { id: '2', text: 'second' } } },
  });

  await waitFor(() => expect(screen.getByTestId('message')).toHaveTextContent('second'));
});

MockSubscriptionLink gives you synchronous control over what the subscription "receives". simulateResult pushes one event through the link immediately — no async waiting required.

Testing Error Handling and Reconnection

GraphQL-level Errors

A subscription can receive a GraphQL error payload (not a transport error). Test that the resolver emits errors correctly:

it('handles GraphQL error in subscription payload', async () => {
  const client = makeClient();
  const errors: unknown[] = [];

  const subscription = client.iterate({
    query: `subscription { messageAdded(roomId: "room-bad") { id text } }`,
  });

  for await (const event of subscription) {
    if (event.errors) {
      errors.push(...event.errors);
      break;
    }
  }

  await client.dispose();

  expect(errors).toHaveLength(1);
  expect(errors[0]).toMatchObject({ message: expect.stringContaining('Unauthorized') });
});

This requires your resolver to throw or return an error for the "bad" room scenario. Configure it in your test schema or resolver mock.

Transport Disconnection and Reconnection

graphql-ws reconnects automatically. Test that your client receives events after a reconnection by forcibly closing the WebSocket mid-stream:

it('resubscribes and continues receiving events after reconnection', async () => {
  const received: unknown[] = [];
  let resolveThree: () => void;
  const threeReceived = new Promise<void>(r => { resolveThree = r; });

  const client = createClient({
    url: `ws://localhost:${port}/graphql`,
    webSocketImpl: WebSocket,
    retryAttempts: 3,
    shouldRetry: () => true,
  });

  const subscription = client.iterate({
    query: `subscription { messageAdded(roomId: "room-reconnect") { id text } }`,
  });

  const iterating = (async () => {
    for await (const event of subscription) {
      received.push(event.data);
      if (received.length === 3) { resolveThree(); break; }
    }
  })();

  // Publish first message, then force-close all WS connections
  await new Promise(r => setTimeout(r, 50));
  pubsub.publish('MESSAGE_ADDED:room-reconnect', {
    messageAdded: { id: '1', text: 'before disconnect' },
  });

  await new Promise(r => setTimeout(r, 50));

  // Forcibly terminate all active WebSocket connections
  wsServer.clients.forEach(ws => ws.terminate());

  // Give client time to reconnect
  await new Promise(r => setTimeout(r, 200));

  pubsub.publish('MESSAGE_ADDED:room-reconnect', {
    messageAdded: { id: '2', text: 'after reconnect' },
  });
  pubsub.publish('MESSAGE_ADDED:room-reconnect', {
    messageAdded: { id: '3', text: 'still after reconnect' },
  });

  await threeReceived;
  await iterating;
  await client.dispose();

  expect(received).toHaveLength(3);
  expect((received[0] as any).messageAdded.text).toBe('before disconnect');
  expect((received[1] as any).messageAdded.text).toBe('after reconnect');
});

Important: ws.terminate() drops the TCP connection without a WebSocket close handshake, which is a harder failure than a clean close. This is what happens in production (network drop, server crash). ws.close() sends a proper close frame and the client may not retry — use terminate() to test reconnection.

Testing Retry Exhaustion

Test that the client stops retrying after the configured limit and the subscription errors out:

it('throws after retryAttempts exhausted', async () => {
  const client = createClient({
    url: `ws://localhost:${port}/graphql`,
    webSocketImpl: WebSocket,
    retryAttempts: 2,
    retryWait: async () => { /* no wait */ },
    shouldRetry: () => true,
  });

  // Shut down the server to prevent any connection
  wsServer.close();

  await expect(async () => {
    for await (const _ of client.iterate({
      query: `subscription { messageAdded(roomId: "x") { id text } }`,
    })) { /* consume */ }
  }).rejects.toThrow();

  // Restore wsServer for other tests — or use beforeEach/afterEach scoping
  await client.dispose();
});

Scope destructive tests (server close, terminate) in their own describe block with beforeAll/afterAll that create isolated server instances. Sharing a server across tests that forcibly close connections causes flakiness.

Structuring Your Test Suite

Organize subscription tests in three layers, mirrored in your directory structure:

test/
  unit/
    resolvers/
      chat.subscription.test.ts   # async generator logic, no server
  integration/
    subscriptions/
      chat.test.ts                # real server, graphql-ws client
      reconnection.test.ts        # transport error scenarios
  components/
    ChatRoom.test.tsx             # MockSubscriptionLink, React Testing Library

Run unit tests on every save (fast, no I/O). Run integration tests in CI and locally before pushing. This separation keeps feedback loops tight.

Set a test timeout appropriate for integration tests — the default 5 seconds is often too short for reconnection tests:

// jest.config.ts
export default {
  testTimeout: 15000, // integration tests can take longer
  // ...
};

// Or per-file:
jest.setTimeout(15000);

Key Takeaways

  • Test the resolver in isolation first. An async generator with a mock iterator is a pure unit test — no server, no timing issues, fast.
  • Use client.iterate() in integration tests. It returns an AsyncIterableIterator that maps cleanly to for await...of. Cleaner than callbacks.
  • Always close iterators. Every test that starts a for await loop must have a path that exits it — break after N items, mock.close(), or client.dispose(). Unclosed iterators cause test timeouts.
  • ws.terminate() for reconnection tests. Clean close() may not trigger retry logic. terminate() simulates a real network failure.
  • MockSubscriptionLink for React tests. Don't stand up a server for component tests. simulateResult() gives synchronous control over the subscription event stream.
  • Isolate destructive tests. Tests that kill WebSocket connections or shut down servers need their own server instance. Sharing state across tests that mutate the transport layer causes non-deterministic failures.
  • Match your integration test timeout to your reconnect delay. If retryWait has any backoff, your test timeout must be larger than the total retry duration.

Subscriptions aren't fundamentally harder than queries — they just have more moving parts. Peel them apart: resolver logic, transport layer, React integration. Test each independently, then add one integration test for the happy path and one for the failure path. That coverage will catch the bugs that matter.

Read more

Start now free