diff --git a/packages/api-graphql/__tests__/AWSAppSyncEventProvider.test.ts b/packages/api-graphql/__tests__/AWSAppSyncEventProvider.test.ts index fbb865be862..a8c29a63e87 100644 --- a/packages/api-graphql/__tests__/AWSAppSyncEventProvider.test.ts +++ b/packages/api-graphql/__tests__/AWSAppSyncEventProvider.test.ts @@ -205,6 +205,43 @@ describe('AppSyncEventProvider', () => { event: JSON.parse(event), }); }); + + test('socket is disconnected after .close() is called', async () => { + expect.assertions(2); + const mockNext = jest.fn(); + + const observer = provider.subscribe({ + appSyncGraphqlEndpoint: 'ws://localhost:8080', + }); + + const event = JSON.stringify({ some: 'data' }); + + observer.subscribe({ + next: mockNext, + error: () => {}, + }); + + await fakeWebSocketInterface?.standardConnectionHandshake(); + await fakeWebSocketInterface?.startAckMessage({ + connectionTimeoutMs: 100, + }); + await fakeWebSocketInterface?.sendDataMessage({ + id: fakeWebSocketInterface?.webSocket.subscriptionId, + type: MESSAGE_TYPES.DATA, + event, + }); + + // events callback returns entire message contents + expect(mockNext).toHaveBeenCalledWith({ + id: fakeWebSocketInterface?.webSocket.subscriptionId, + type: MESSAGE_TYPES.DATA, + event: JSON.parse(event), + }); + + await provider.close(); + + expect(fakeWebSocketInterface.hasClosed).resolves.toBeUndefined(); + }); }); }); }); diff --git a/packages/api-graphql/__tests__/helpers.ts b/packages/api-graphql/__tests__/helpers.ts index ac26232dd90..424fc4e3129 100644 --- a/packages/api-graphql/__tests__/helpers.ts +++ b/packages/api-graphql/__tests__/helpers.ts @@ -258,9 +258,8 @@ export class FakeWebSocketInterface { /** * Run a command and resolve to allow internal behavior to execute */ - async runAndResolve(fn) { + async runAndResolve(fn: Function) { await fn(); - await Promise.resolve(); } /** @@ -310,6 +309,10 @@ class FakeWebSocket implements WebSocket { close(code?: number, reason?: string): void { const closeResolver = this.closeResolverFcn(); if (closeResolver) closeResolver(Promise.resolve(undefined)); + + try { + this.onclose(new CloseEvent('', {})); + } catch {} } send(data: string | ArrayBufferLike | Blob | ArrayBufferView): void { const parsedInput = JSON.parse(String(data)); diff --git a/packages/api-graphql/src/Providers/AWSAppSyncEventsProvider/index.ts b/packages/api-graphql/src/Providers/AWSAppSyncEventsProvider/index.ts index eff48289acc..859f69fcde6 100644 --- a/packages/api-graphql/src/Providers/AWSAppSyncEventsProvider/index.ts +++ b/packages/api-graphql/src/Providers/AWSAppSyncEventsProvider/index.ts @@ -53,6 +53,27 @@ export class AWSAppSyncEventProvider extends AWSWebSocketProvider { return PROVIDER_NAME; } + close() { + return new Promise((resolve, reject) => { + super.close(); + if (this.awsRealTimeSocket) { + this.awsRealTimeSocket.onclose = (_: CloseEvent) => { + this.subscriptionObserverMap = new Map(); + this.awsRealTimeSocket = undefined; + resolve(); + }; + + this.awsRealTimeSocket.onerror = (err: any) => { + reject(err); + }; + + this.awsRealTimeSocket.close(); + } else { + resolve(); + } + }); + } + public async connect(options: AWSAppSyncEventProviderOptions) { super.connect(options); } diff --git a/packages/api-graphql/src/Providers/AWSWebSocketProvider/index.ts b/packages/api-graphql/src/Providers/AWSWebSocketProvider/index.ts index 3553a008123..4a5e0d8f3b0 100644 --- a/packages/api-graphql/src/Providers/AWSWebSocketProvider/index.ts +++ b/packages/api-graphql/src/Providers/AWSWebSocketProvider/index.ts @@ -80,7 +80,7 @@ export abstract class AWSWebSocketProvider { protected logger: ConsoleLogger; protected subscriptionObserverMap = new Map(); - private awsRealTimeSocket?: WebSocket; + protected awsRealTimeSocket?: WebSocket; private socketStatus: SOCKET_STATUS = SOCKET_STATUS.CLOSED; private keepAliveTimeoutId?: ReturnType; private keepAliveTimeout = DEFAULT_KEEP_ALIVE_TIMEOUT; diff --git a/packages/api-graphql/src/internals/events/index.ts b/packages/api-graphql/src/internals/events/index.ts index 49860a92049..c1101111fd4 100644 --- a/packages/api-graphql/src/internals/events/index.ts +++ b/packages/api-graphql/src/internals/events/index.ts @@ -154,8 +154,19 @@ async function post( } } -function closeAll(): void { - eventProvider.close(); +/** + * @experimental API may change in future versions + * + * Close WebSocket connection, disconnect listeners and reconnect observers + * + * @example + * await events.closeAll() + * + * @returns void on success + * @throws on error + */ +async function closeAll(): Promise { + await eventProvider.close(); } export { connect, post, closeAll };