From 985850f2e1f506220695c1f479c4e6cae6419a9b Mon Sep 17 00:00:00 2001
From: Laurens Lavaert <llavaert@gmail.com>
Date: Mon, 16 Sep 2019 10:41:20 +0200
Subject: [PATCH] Add subscription support for testing client & improved
 typings

---
 .../src/createTestClient.ts                   | 85 +++++++++++++------
 1 file changed, 58 insertions(+), 27 deletions(-)

diff --git a/packages/apollo-server-testing/src/createTestClient.ts b/packages/apollo-server-testing/src/createTestClient.ts
index 44362f2b1a9..d4a6d479138 100644
--- a/packages/apollo-server-testing/src/createTestClient.ts
+++ b/packages/apollo-server-testing/src/createTestClient.ts
@@ -1,43 +1,74 @@
 import { ApolloServerBase } from 'apollo-server-core';
 import { GraphQLResponse } from 'apollo-server-types';
-import { print, DocumentNode } from 'graphql';
+import { print, DocumentNode, subscribe, GraphQLSchema, GraphQLError, ExecutionResult } from 'graphql';
 
 type StringOrAst = string | DocumentNode;
 
-// A query must not come with a mutation (and vice versa).
-type Query = {
+interface Query<T = any> {
   query: StringOrAst;
-  mutation?: undefined;
-  variables?: {
-    [name: string]: any;
-  };
-  operationName?: string;
-};
-type Mutation = {
+  variables?: T;
+}
+
+interface Mutation<T = any> {
   mutation: StringOrAst;
-  query?: undefined;
-  variables?: {
-    [name: string]: any;
-  };
-  operationName?: string;
-};
+  variables?: T;
+}
+
+interface Subscription<T = any> {
+  subscription: StringOrAst;
+  variables?: T;
+}
 
 export interface ApolloServerTestClient {
-  query: (query: Query) => Promise<GraphQLResponse>;
-  mutate: (mutation: Mutation) => Promise<GraphQLResponse>;
+  query: <T = any>(query: Query<T>) => Promise<GraphQLResponse>;
+  mutate: <T = any>(mutation: Mutation<T>) => Promise<GraphQLResponse>;
+  subscription: <T = any>(subscription: Subscription<T>) => Promise<any>;
 }
 
+const isAsyncIterable = (iterator: any): iterator is AsyncIterableIterator<any> => {
+  return typeof iterator[Symbol.asyncIterator] === 'function';
+};
+
+const sleep = (ms: number) => {
+  if (ms <= 0) return false;
+  return new Promise(resolve => {
+    setTimeout(resolve, ms);
+  });
+};
+
+export const waitForSubscription = async (
+  server: ApolloServerBase,
+  subscription: DocumentNode,
+  variables?: any
+): Promise<any | ReadonlyArray<GraphQLError>> => {
+  // @ts-ignore
+  const schema: GraphQLSchema = server.schema;
+  // @ts-ignore
+  const context = server.context();
+  const iterator = await subscribe(schema, subscription, {}, context, variables);
+  await sleep(5); // Else subscription doesn't apply
+  return () =>
+    new Promise<any | ReadonlyArray<GraphQLError>>(async (resolve, reject) => {
+      if ('errors' in iterator) {
+        reject(iterator.errors);
+      }
+      if (isAsyncIterable(iterator)) {
+        for await (const next of iterator) {
+          const val: ExecutionResult<any> = next;
+          if (val.errors) reject(val.errors);
+          const result = val.data;
+          if (result) resolve(result);
+        }
+      }
+    });
+};
+
 export default (server: ApolloServerBase): ApolloServerTestClient => {
   const executeOperation = server.executeOperation.bind(server);
-  const test = ({ query, mutation, ...args }: Query | Mutation) => {
+  const test = ({ query, mutation, subscription, ...args }: Query & Mutation & Subscription) => {
+    if (subscription) return waitForSubscription(server, subscription as DocumentNode, args.variables);
     const operation = query || mutation;
-
-    if (!operation || (query && mutation)) {
-      throw new Error(
-        'Either `query` or `mutation` must be passed, but not both.',
-      );
-    }
-
+    if (!operation) throw new Error('Either `query` or `mutation` or `subscription` must be passed.');
     return executeOperation({
       // Convert ASTs, which are produced by `graphql-tag` but not currently
       // used by `executeOperation`, to a String using `graphql/language/print`.
@@ -46,5 +77,5 @@ export default (server: ApolloServerBase): ApolloServerTestClient => {
     });
   };
 
-  return { query: test, mutate: test };
+  return { query: test, mutate: test, subscription: test };
 };