fix(api-graphql): close WS on events.closeAll (#13989)

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();
+				});
 			});
 		});
 	});

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));

packages/api-graphql/src/Providers/AWSAppSyncEventsProvider/index.ts

@@ -53,6 +53,27 @@ export class AWSAppSyncEventProvider extends AWSWebSocketProvider {
 		return PROVIDER_NAME;
 	}
 
+	close() {
+		return new Promise<void>((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);
 	}

packages/api-graphql/src/Providers/AWSWebSocketProvider/index.ts

@@ -80,7 +80,7 @@ export abstract class AWSWebSocketProvider {
 	protected logger: ConsoleLogger;
 	protected subscriptionObserverMap = new Map<string, ObserverQuery>();
 
-	private awsRealTimeSocket?: WebSocket;
+	protected awsRealTimeSocket?: WebSocket;
 	private socketStatus: SOCKET_STATUS = SOCKET_STATUS.CLOSED;
 	private keepAliveTimeoutId?: ReturnType<typeof setTimeout>;
 	private keepAliveTimeout = DEFAULT_KEEP_ALIVE_TIMEOUT;

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<void> {
+	await eventProvider.close();
 }
 
 export { connect, post, closeAll };