feat(client): Async iterator for subscriptions (#486)

src/__tests__/client.ts

@@ -14,6 +14,7 @@ import {
 } from '../common';
 import { startRawServer, startWSTServer as startTServer } from './utils';
 import { ExecutionResult } from 'graphql';
+import { pong } from './fixtures/simple';
 
 // silence console.error calls for nicer tests overview
 const consoleError = console.error;
@@ -2324,3 +2325,261 @@ describe('events', () => {
     expect(expected).toBeCalledTimes(6);
   });
 });
+
+describe('iterate', () => {
+  it('should iterate a single result query', async () => {
+    const { url } = await startTServer();
+
+    const client = createClient({
+      url,
+      retryAttempts: 0,
+      onNonLazyError: noop,
+    });
+
+    const iterator = client.iterate({
+      query: '{ getValue }',
+    });
+
+    await expect(iterator.next()).resolves.toMatchInlineSnapshot(`
+      {
+        "done": false,
+        "value": {
+          "data": {
+            "getValue": "value",
+          },
+        },
+      }
+    `);
+
+    await expect(iterator.next()).resolves.toMatchInlineSnapshot(`
+      {
+        "done": true,
+        "value": undefined,
+      }
+    `);
+  });
+
+  it('should iterate over subscription events', async () => {
+    const { url } = await startTServer();
+
+    const client = createClient({
+      url,
+      retryAttempts: 0,
+      onNonLazyError: noop,
+    });
+
+    const iterator = client.iterate({
+      query: 'subscription { greetings }',
+    });
+
+    // Hi
+    await expect(iterator.next()).resolves.toBeDefined();
+    // Bonjour
+    await expect(iterator.next()).resolves.toBeDefined();
+    // Hola
+    await expect(iterator.next()).resolves.toBeDefined();
+    // Ciao
+    await expect(iterator.next()).resolves.toBeDefined();
+    // Zdravo
+    await expect(iterator.next()).resolves.toBeDefined();
+
+    await expect(iterator.next()).resolves.toMatchInlineSnapshot(`
+      {
+        "done": true,
+        "value": undefined,
+      }
+    `);
+  });
+
+  it('should report execution errors to iterator', async () => {
+    const { url } = await startTServer();
+
+    const client = createClient({
+      url,
+      retryAttempts: 0,
+      onNonLazyError: noop,
+    });
+
+    const iterator = client.iterate({
+      query: 'subscription { throwing }',
+    });
+
+    await expect(iterator.next()).resolves.toMatchInlineSnapshot(`
+      {
+        "done": false,
+        "value": {
+          "errors": [
+            {
+              "locations": [
+                {
+                  "column": 16,
+                  "line": 1,
+                },
+              ],
+              "message": "Kaboom!",
+              "path": [
+                "throwing",
+              ],
+            },
+          ],
+        },
+      }
+    `);
+
+    await expect(iterator.next()).resolves.toMatchInlineSnapshot(`
+      {
+        "done": true,
+        "value": undefined,
+      }
+    `);
+  });
+
+  it('should throw in iterator connection errors', async () => {
+    const { url, ...server } = await startTServer();
+
+    const client = createClient({
+      url,
+      retryAttempts: 0,
+      onNonLazyError: noop,
+    });
+
+    const pingKey = Math.random().toString();
+    const iterator = client.iterate({
+      query: `subscription { ping(key: "${pingKey}") }`,
+    });
+
+    pong(pingKey);
+    await expect(iterator.next()).resolves.toMatchInlineSnapshot(`
+      {
+        "done": false,
+        "value": {
+          "data": {
+            "ping": "pong",
+          },
+        },
+      }
+    `);
+
+    await server.dispose(false);
+
+    await expect(iterator.next()).rejects.toEqual(
+      // forceful close
+      expect.objectContaining({
+        code: 1006,
+        reason: '',
+      }),
+    );
+  });
+
+  it('should complete subscription when iterator loop breaks', async () => {
+    const { url, ...server } = await startTServer();
+
+    const client = createClient({
+      url,
+      retryAttempts: 0,
+      onNonLazyError: noop,
+    });
+
+    const pingKey = Math.random().toString();
+    const iterator = client.iterate({
+      query: `subscription { ping(key: "${pingKey}") }`,
+    });
+    iterator.return = jest.fn(iterator.return);
+
+    setTimeout(() => pong(pingKey), 0);
+
+    for await (const val of iterator) {
+      expect(val).toMatchInlineSnapshot(`
+        {
+          "data": {
+            "ping": "pong",
+          },
+        }
+      `);
+      break;
+    }
+
+    expect(iterator.return).toHaveBeenCalled();
+
+    await server.waitForClientClose();
+  });
+
+  it('should complete subscription when iterator loop throws', async () => {
+    const { url, ...server } = await startTServer();
+
+    const client = createClient({
+      url,
+      retryAttempts: 0,
+      onNonLazyError: noop,
+    });
+
+    const pingKey = Math.random().toString();
+    const iterator = client.iterate({
+      query: `subscription { ping(key: "${pingKey}") }`,
+    });
+    iterator.return = jest.fn(iterator.return);
+
+    setTimeout(() => pong(pingKey), 0);
+
+    await expect(async () => {
+      for await (const val of iterator) {
+        expect(val).toMatchInlineSnapshot(`
+          {
+            "data": {
+              "ping": "pong",
+            },
+          }
+        `);
+        throw new Error(':)');
+      }
+    }).rejects.toBeDefined();
+
+    expect(iterator.return).toHaveBeenCalled();
+
+    await server.waitForClientClose();
+  });
+
+  it('should complete subscription when calling return directly on iterator', async () => {
+    const { url, ...server } = await startTServer();
+
+    const client = createClient({
+      url,
+      retryAttempts: 0,
+      onNonLazyError: noop,
+    });
+
+    const pingKey = Math.random().toString();
+    const iterator = client.iterate({
+      query: `subscription { ping(key: "${pingKey}") }`,
+    });
+
+    pong(pingKey);
+
+    await expect(iterator.next()).resolves.toMatchInlineSnapshot(`
+      {
+        "done": false,
+        "value": {
+          "data": {
+            "ping": "pong",
+          },
+        },
+      }
+    `);
+
+    await expect(iterator.return?.()).resolves.toMatchInlineSnapshot(`
+      {
+        "done": true,
+        "value": undefined,
+      }
+    `);
+
+    await expect(iterator.next()).resolves.toMatchInlineSnapshot(`
+      {
+        "done": true,
+        "value": undefined,
+      }
+    `);
+
+    await server.waitForClientClose();
+  });
+});

src/__tests__/fixtures/simple.ts

@@ -100,6 +100,12 @@ export const schemaConfig: GraphQLSchemaConfig = {
           };
         },
       },
+      throwing: {
+        type: new GraphQLNonNull(GraphQLString),
+        subscribe: async function () {
+          throw new Error('Kaboom!');
+        },
+      },
     },
   }),
 };

src/client.ts

@@ -440,6 +440,13 @@ export interface Client extends Disposable {
     payload: SubscribePayload,
     sink: Sink<ExecutionResult<Data, Extensions>>,
   ): () => void;
+  /**
+   * Subscribes and iterates over emitted results from the WebSocket
+   * through the returned async iterator.
+   */
+  iterate<Data = Record<string, unknown>, Extensions = unknown>(
+    payload: SubscribePayload,
+  ): AsyncIterableIterator<ExecutionResult<Data, Extensions>>;
   /**
    * Terminates the WebSocket abruptly and immediately.
    *
@@ -972,6 +979,73 @@ export function createClient<
         if (!done) releaser();
       };
     },
+    iterate(request) {
+      const pending: ExecutionResult<
+        // TODO: how to not use `any` and not have a redundant function signature?
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        any,
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        any
+      >[] = [];
+      const deferred = {
+        done: false,
+        error: null as unknown,
+        resolve: () => {
+          // noop
+        },
+      };
+      const dispose = this.subscribe(request, {
+        next(val) {
+          pending.push(val);
+          deferred.resolve();
+        },
+        error(err) {
+          deferred.done = true;
+          deferred.error = err;
+          deferred.resolve();
+        },
+        complete() {
+          deferred.done = true;
+          deferred.resolve();
+        },
+      });
+
+      const iterator = (async function* iterator() {
+        for (;;) {
+          if (!pending.length) {
+            // only wait if there are no pending messages available
+            await new Promise<void>((resolve) => (deferred.resolve = resolve));
+          }
+          // first flush
+          while (pending.length) {
+            // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+            yield pending.shift()!;
+          }
+          // then error
+          if (deferred.error) {
+            throw deferred.error;
+          }
+          // or complete
+          if (deferred.done) {
+            return;
+          }
+        }
+      })();
+      iterator.throw = async (err) => {
+        if (!deferred.done) {
+          deferred.done = true;
+          deferred.error = err;
+          deferred.resolve();
+        }
+        return { done: true, value: undefined };
+      };
+      iterator.return = async () => {
+        dispose();
+        return { done: true, value: undefined };
+      };
+
+      return iterator;
+    },
     async dispose() {
       disposed = true;
       if (connecting) {