Fix onHealthCheck for lambda failing to correctly handle returned pro…

…mise (#4969)

* Ensure lambda healthcheck callback waits on returned promise

This fixes the issue "Custom onHealthCheck function is intermittently called" reported at #3999.
This fix ensures the healthcheck handling code does not carry on executing and waits for the returned promise to be resolved or rejected.

* Add new tests to exercise lambda healthcheck

The tests validates that a healthcheck endpoint exists and correctly handles promises returned from the callback.

This commit also moves the 'mock server'-like code (responsible for transforming lambda response into HTTP response) has been moved to a common location where it is reused in the newly created apollo-server-lambda/src/ApolloServer.test.ts.

packages/apollo-server-lambda/src/ApolloServer.ts

@@ -197,6 +197,7 @@ export class ApolloServer extends ApolloServerBase {
                 },
               });
             });
+          return;
         } else {
           return callback(null, successfulResponse);
         }

packages/apollo-server-lambda/src/__tests__/ApolloServer.test.ts

@@ -0,0 +1,87 @@
+import request from 'supertest';
+import {createMockServer} from './mockServer';
+import { gql } from 'apollo-server-core';
+import {
+  ApolloServer,
+  CreateHandlerOptions
+} from '../ApolloServer';
+
+const typeDefs = gql`
+  type Query {
+    hello: String
+  }
+`;
+
+const resolvers = {
+  Query: {
+    hello: () => 'hi',
+  },
+};
+
+describe('apollo-server-lambda', () => {
+  const createLambda = (
+    options: Partial<CreateHandlerOptions> = {},
+  ) => {
+    const server = new ApolloServer({
+      typeDefs,
+      resolvers
+    });
+    const handler = server.createHandler(options);
+    return createMockServer(handler);
+  }
+
+  describe('healthchecks', () => {
+
+    it('creates a healthcheck endpoint', async () => {
+      const app = createLambda();
+
+      const req = request(app)
+        .get('/.well-known/apollo/server-health');
+
+      return req.then((res: any) => {
+        expect(res.status).toEqual(200);
+        expect(res.body).toEqual({ status: 'pass' });
+        expect(res.headers['content-type']).toEqual('application/json');
+      });
+    });
+
+    it('provides a callback for the healthcheck', async () => {
+      const app = createLambda({
+        onHealthCheck: async () => {
+          return new Promise((resolve) => {
+            return resolve("Success!");
+          });
+        }
+      });
+
+      const req = request(app)
+        .get('/.well-known/apollo/server-health');
+
+      return req.then((res: any) => {
+        expect(res.status).toEqual(200);
+        expect(res.body).toEqual({ status: 'pass' });
+        expect(res.headers['content-type']).toEqual('application/json');
+      });
+    });
+
+    it('returns a 503 if healthcheck fails', async () => {
+      const app = createLambda({
+        onHealthCheck: async () => {
+          return new Promise(() => {
+            throw new Error("Failed to connect!");
+          });
+        }
+      });
+
+      const req = request(app)
+        .get('/.well-known/apollo/server-health');
+
+      return req.then((res: any) => {
+        expect(res.status).toEqual(503);
+        expect(res.body).toEqual({ status: 'fail' });
+        expect(res.headers['content-type']).toEqual('application/json');
+      });
+    });
+  });
+
+});

packages/apollo-server-lambda/src/__tests__/lambdaApollo.test.ts

@@ -5,12 +5,9 @@ import testSuite, {
   NODE_MAJOR_VERSION,
 } from 'apollo-server-integration-testsuite';
 import { Config } from 'apollo-server-core';
-import { IncomingMessage, ServerResponse } from 'http';
-
-import url from 'url';
 import gql from 'graphql-tag';
 import request from 'supertest';
-import { APIGatewayProxyCallback } from "aws-lambda";
+import {createMockServer} from './mockServer';
 
 const createLambda = (options: CreateAppOptions = {}) => {
   const server = new ApolloServer(
@@ -19,52 +16,8 @@ const createLambda = (options: CreateAppOptions = {}) => {
 
   const handler = server.createHandler();
 
-  return (req: IncomingMessage, res: ServerResponse) => {
-    // return 404 if path is /bogus-route to pass the test, lambda doesn't have paths
-    if (req.url && req.url.includes('/bogus-route')) {
-      res.statusCode = 404;
-      return res.end();
-    }
-
-    let body = '';
-    req.on('data', chunk => (body += chunk));
-    req.on('end', () => {
-      const urlObject = url.parse(req.url || '', true);
-      const event = {
-        httpMethod: req.method,
-        body: body,
-        path: req.url,
-        queryStringParameters: urlObject.query,
-        requestContext: {
-          path: urlObject.pathname,
-        },
-        headers: req.headers,
-      };
-      const callback: APIGatewayProxyCallback = (error, result) => {
-        if (error) {
-          throw error;
-        } else {
-          result = result as NonNullable<typeof result>;
-        }
-        res.statusCode = result.statusCode;
-        for (let key in result.headers) {
-          if (result.headers.hasOwnProperty(key)) {
-            if (typeof result.headers[key] === 'boolean') {
-              res.setHeader(key, result.headers[key].toString());
-            } else {
-              // Without casting this to `any`, TS still believes `boolean`
-              // is possible.
-              res.setHeader(key, result.headers[key] as any);
-            }
-          }
-        }
-        res.write(result.body);
-        res.end();
-      };
-      handler(event as any, {} as any, callback);
-    });
-  };
-};
+  return createMockServer(handler);
+}
 
 describe('integration:Lambda', () => {
   testSuite(createLambda);

packages/apollo-server-lambda/src/__tests__/mockServer.ts

@@ -0,0 +1,62 @@
+import url from 'url';
+import { IncomingMessage, ServerResponse } from 'http';
+import {
+  APIGatewayProxyCallback,
+  APIGatewayProxyEvent,
+  Context as LambdaContext,
+  Handler,
+  APIGatewayProxyResult
+} from "aws-lambda";
+
+export function createMockServer(handler: Handler<APIGatewayProxyEvent, APIGatewayProxyResult>) {
+  return (req: IncomingMessage, res: ServerResponse) => {
+    // return 404 if path is /bogus-route to pass the test, lambda doesn't have paths
+    if (req.url && req.url.includes('/bogus-route')) {
+      res.statusCode = 404;
+      return res.end();
+    }
+
+    let body = '';
+    req.on('data', chunk => (body += chunk));
+    req.on('end', () => {
+      const urlObject = url.parse(req.url || '', true);
+      const event = {
+        httpMethod: req.method,
+        body: body,
+        path: req.url,
+        queryStringParameters: urlObject.query,
+        requestContext: {
+          path: urlObject.pathname,
+        },
+        headers: req.headers,
+      };
+      const callback: APIGatewayProxyCallback = (error, result) => {
+        if (error) {
+          throw error;
+        } else {
+          result = result as NonNullable<typeof result>;
+        }
+        res.statusCode = result.statusCode;
+        for (let key in result.headers) {
+          if (result.headers.hasOwnProperty(key)) {
+            if (typeof result.headers[key] === 'boolean') {
+              res.setHeader(key, result.headers[key].toString());
+            } else {
+              // Without casting this to `any`, TS still believes `boolean`
+              // is possible.
+              res.setHeader(key, result.headers[key] as any);
+            }
+          }
+        }
+        res.write(result.body);
+        res.end();
+      };
+
+      handler(
+        event as APIGatewayProxyEvent,
+        {} as LambdaContext,
+        callback as APIGatewayProxyCallback
+      );
+    });
+  };
+};