From 565d96e37defabd6ba8851d198bbc7d4a4e8fb22 Mon Sep 17 00:00:00 2001
From: Trevor Scheer <trevor.scheer@gmail.com>
Date: Wed, 1 Apr 2020 20:43:50 -0700
Subject: [PATCH] Unbreak the onSchemaChange contract

---
 .../integration/networkRequests.test.ts       | 40 ++++++++++++++-----
 packages/apollo-gateway/src/index.ts          | 24 +++++------
 2 files changed, 41 insertions(+), 23 deletions(-)

diff --git a/packages/apollo-gateway/src/__tests__/integration/networkRequests.test.ts b/packages/apollo-gateway/src/__tests__/integration/networkRequests.test.ts
index bd15db3c92d..df976f8029a 100644
--- a/packages/apollo-gateway/src/__tests__/integration/networkRequests.test.ts
+++ b/packages/apollo-gateway/src/__tests__/integration/networkRequests.test.ts
@@ -102,7 +102,6 @@ afterEach(() => {
   expect(nock.isDone()).toBeTruthy();
   nock.cleanAll();
   nock.restore();
-  jest.useRealTimers();
 });
 
 it('Queries remote endpoints for their SDLs', async () => {
@@ -399,21 +398,44 @@ describe('Downstream service health checks', () => {
       let resolve: () => void;
       const schemaChangeBlocker = new Promise(res => (resolve = res));
 
-      const onChange = jest
-        .fn()
-        .mockImplementationOnce(() => resolve())
-
       const gateway = new ApolloGateway({ serviceHealthCheck: true, logger });
       // @ts-ignore for testing purposes, a short pollInterval is ideal so we'll override here
-      gateway.experimental_pollInterval = 300;
+      gateway.experimental_pollInterval = 100;
 
+      // load the gateway as usual
       await gateway.load({ engine: { apiKeyHash, graphId } });
-      gateway.onSchemaChange(onChange);
+      expect(gateway.schema!.getType('User')!.description).toBe('This is my User');
+
+      // @ts-ignore for testing purposes, we'll call the original `updateComposition`
+      // function from our mock
+      const original = gateway.updateComposition;
+      const mockUpdateComposition = jest
+        .fn(original)
+        .mockImplementationOnce(async opts => {
+          // mock the first poll and handle the error which would otherwise be caught
+          // and logged from within the `pollServices` class method
+          await expect(original.apply(gateway, [opts]))
+            .rejects
+            .toThrowErrorMatchingInlineSnapshot(
+              `"500: Internal Server Error"`,
+            );
+          // finally resolve the promise which drives this test
+          resolve();
+        });
+
+      // @ts-ignore for testing purposes, replace the `updateComposition`
+      // function on the gateway with our mock
+      gateway.updateComposition = mockUpdateComposition;
+
+      // This kicks off polling within the gateway
+      gateway.onSchemaChange(() => {});
 
       await schemaChangeBlocker;
-      expect(onChange.mock.calls.length).toBe(1);
 
+      // At this point, the mock update should have been called but the schema
+      // should not have updated to the new one.
+      expect(mockUpdateComposition.mock.calls.length).toBe(1);
       expect(gateway.schema!.getType('User')!.description).toBe('This is my User');
     });
-  })
+  });
 });
diff --git a/packages/apollo-gateway/src/index.ts b/packages/apollo-gateway/src/index.ts
index e10dfdcdeb2..edfa1e38e5e 100644
--- a/packages/apollo-gateway/src/index.ts
+++ b/packages/apollo-gateway/src/index.ts
@@ -379,8 +379,6 @@ export class ApolloGateway implements GraphQLService {
           'The gateway did not update its schema due to failed service health checks.  ' +
           'The gateway will continue to operate with the previous schema and reattempt updates.' + e
         );
-        // TODO: Questionable notify here, but it does enable code-driven testing via schema listener callbacks
-        this.notifySchemaListeners(this.schema);
         throw e;
       }
     }
@@ -391,7 +389,15 @@ export class ApolloGateway implements GraphQLService {
     if (this.queryPlanStore) this.queryPlanStore.flush();
 
     this.schema = this.createSchema(result.serviceDefinitions);
-    this.notifySchemaListeners(this.schema);
+
+    // Notify the schema listeners of the updated schema
+    try {
+      this.onSchemaChangeListeners.forEach(listener => listener(this.schema!));
+    } catch (e) {
+      this.logger.error(
+        "An error was thrown from an 'onSchemaChange' listener. " +
+        "The schema will still update: " + (e && e.message || e));
+    }
 
     if (this.experimental_didUpdateComposition) {
       this.experimental_didUpdateComposition(
@@ -414,16 +420,6 @@ export class ApolloGateway implements GraphQLService {
     }
   }
 
-  private notifySchemaListeners(schema?: GraphQLSchema) {
-    try {
-      this.onSchemaChangeListeners.forEach(listener => listener(schema!));
-    } catch (e) {
-      this.logger.error(
-        "An error was thrown from an 'onSchemaChange' listener. " +
-        "The schema will still update: " + (e && e.message || e));
-    }
-  }
-
   /**
    * This can be used without an argument in order to perform an ad-hoc health check
    * of the downstream services like so:
@@ -780,7 +776,7 @@ export class ApolloGateway implements GraphQLService {
 
   public async stop() {
     if (this.pollingTimer) {
-      clearInterval(this.pollingTimer);
+      clearTimeout(this.pollingTimer);
       this.pollingTimer = undefined;
     }
   }