From 5b580a0e946ad1bdd8c17c1e6e4bcbc305449d08 Mon Sep 17 00:00:00 2001
From: Francis Gulotta <wizard@roborooter.com>
Date: Mon, 16 Aug 2021 00:12:43 -0400
Subject: [PATCH] fix: gatwayHandler types are now more narrow

---
 lib/gateway.ts                                |  4 +--
 lib/index-test.ts                             | 25 ++++++++++++++
 lib/messages/subscribe-test.ts                |  2 +-
 .../{mockServerContext.ts => mockServer.ts}   | 34 ++++++++++++-------
 lib/types.ts                                  |  6 ++--
 5 files changed, 53 insertions(+), 18 deletions(-)
 create mode 100644 lib/index-test.ts
 rename lib/test/{mockServerContext.ts => mockServer.ts} (51%)

diff --git a/lib/gateway.ts b/lib/gateway.ts
index 27634d30..e30ce900 100644
--- a/lib/gateway.ts
+++ b/lib/gateway.ts
@@ -1,6 +1,6 @@
-import { Handler } from 'aws-lambda'
 import { GRAPHQL_TRANSPORT_WS_PROTOCOL, MessageType } from 'graphql-ws'
 import {
+  ApiGatewayHandler,
   APIGatewayWebSocketEvent,
   ServerClosure,
   WebsocketResponse,
@@ -13,7 +13,7 @@ import { connection_init } from './messages/connection_init'
 import { pong } from './messages/pong'
 
 export const handleGatewayEvent =
-  (server: ServerClosure): Handler<APIGatewayWebSocketEvent, WebsocketResponse> =>
+  (server: ServerClosure): ApiGatewayHandler<APIGatewayWebSocketEvent, WebsocketResponse> =>
     async (event) => {
       if (!event.requestContext) {
         return {
diff --git a/lib/index-test.ts b/lib/index-test.ts
new file mode 100644
index 00000000..0ae991a9
--- /dev/null
+++ b/lib/index-test.ts
@@ -0,0 +1,25 @@
+import { assert } from 'chai'
+import { Handler } from 'aws-lambda'
+import { tables } from '@architect/sandbox'
+import { createInstance } from '.'
+import { mockServerArgs } from './test/mockServer'
+import { APIGatewayWebSocketEvent, WebsocketResponse } from './types'
+
+describe('createInstance', () => {
+  describe('gatewayHandler', () => {
+    before(async () => {
+      await tables.start({ cwd: './mocks/arc-basic-events', quiet: true })
+    })
+
+    after(async () => {
+      tables.end()
+    })
+
+    it('is type compatible with aws-lambda handler', async () => {
+      const server = createInstance(await mockServerArgs())
+
+      const gatewayHandler: Handler<APIGatewayWebSocketEvent, WebsocketResponse> = server.gatewayHandler
+      assert.ok(gatewayHandler)
+    })
+  })
+})
diff --git a/lib/messages/subscribe-test.ts b/lib/messages/subscribe-test.ts
index d963b2f2..6c3241dc 100644
--- a/lib/messages/subscribe-test.ts
+++ b/lib/messages/subscribe-test.ts
@@ -1,7 +1,7 @@
 import { assert } from 'chai'
 import { tables } from '@architect/sandbox'
 import { subscribe } from './subscribe'
-import { mockServerContext } from '../test/mockServerContext'
+import { mockServerContext } from '../test/mockServer'
 import { connection_init } from './connection_init'
 import { equals } from '@aws/dynamodb-expressions'
 import { collect } from 'streaming-iterables'
diff --git a/lib/test/mockServerContext.ts b/lib/test/mockServer.ts
similarity index 51%
rename from lib/test/mockServerContext.ts
rename to lib/test/mockServer.ts
index ab29df59..8dc9b11f 100644
--- a/lib/test/mockServerContext.ts
+++ b/lib/test/mockServer.ts
@@ -1,3 +1,4 @@
+/* eslint-disable @typescript-eslint/no-empty-function */
 import { makeExecutableSchema } from '@graphql-tools/schema'
 import { tables as arcTables } from '@architect/functions'
 import { makeServerClosure } from '../makeServerClosure'
@@ -33,26 +34,33 @@ const schema = makeExecutableSchema({
   resolvers,
 })
 
+const ensureName = (tables: any, table: string) => {
+  const actualTableName = tables.name(table)
+  if (!actualTableName) {
+    throw new Error(`No table found for ${table}`)
+  }
+  return actualTableName
+}
 
-export const mockServerContext = async (args: Partial<ServerArgs>): Promise<ServerClosure> => {
+export const mockServerArgs = async (args: Partial<ServerArgs> = {}): Promise<ServerArgs> => {
   const tables = await arcTables()
 
-  const ensureName = (table) => {
-    const actualTableName = tables.name(table)
-    if (!actualTableName) {
-      throw new Error(`No table found for ${table}`)
-    }
-    return actualTableName
-  }
-
-  return makeServerClosure({
+  return {
     dynamodb: arcTables.db,
     schema,
     tableNames: {
-      connections: ensureName('Connection'),
-      subscriptions: ensureName('Subscription'),
+      connections: ensureName(tables, 'Connection'),
+      subscriptions: ensureName(tables, 'Subscription'),
+    },
+    apiGatewayManagementApi: {
+      postToConnection: () => ({ promise: async () => { } }),
+      deleteConnection: () => ({ promise: async () => { } }),
     },
     onError: (err) => { console.log('onError'); throw err },
     ...args,
-  })
+  }
+}
+
+export const mockServerContext = async (args?: Partial<ServerArgs>): Promise<ServerClosure> => {
+  return makeServerClosure(await mockServerArgs(args))
 }
diff --git a/lib/types.ts b/lib/types.ts
index 3bee46cb..39aed707 100644
--- a/lib/types.ts
+++ b/lib/types.ts
@@ -2,7 +2,7 @@
 /* eslint-disable @typescript-eslint/ban-types */
 import { ConnectionInitMessage, PingMessage, PongMessage } from 'graphql-ws'
 import { DataMapper } from '@aws/dynamodb-data-mapper'
-import { APIGatewayEventRequestContext, APIGatewayProxyEvent, Handler } from 'aws-lambda'
+import { APIGatewayEventRequestContext, APIGatewayProxyEvent } from 'aws-lambda'
 import { GraphQLResolveInfo, GraphQLSchema } from 'graphql'
 import { DynamoDB } from 'aws-sdk'
 import { Subscription } from './model/Subscription'
@@ -48,7 +48,7 @@ export type ServerClosure = {
 } & Omit<ServerArgs, 'tableNames'>
 
 export interface ServerInstance {
-  gatewayHandler: Handler<APIGatewayWebSocketEvent, WebsocketResponse>
+  gatewayHandler: ApiGatewayHandler<APIGatewayWebSocketEvent, WebsocketResponse>
   stateMachineHandler: (input: StateFunctionInput) => Promise<StateFunctionInput>
   publish: (event: PubSubEvent) => Promise<void>
   complete: (event: PubSubEvent) => Promise<void>
@@ -123,3 +123,5 @@ export interface SubscribeOptions {
   onComplete?: (...args: SubscribeArgs) => void | Promise<void>
   onAfterSubscribe?: (...args: SubscribeArgs) => PubSubEvent | Promise<PubSubEvent> | undefined | Promise<undefined>
 }
+
+export type ApiGatewayHandler<TEvent = any, TResult = any> = (event: TEvent) => void | Promise<TResult>