From d6eede3a68df1e6c49927e62ce0b73d6db615a8a Mon Sep 17 00:00:00 2001
From: Ivan Artemiev <29709626+iartemiev@users.noreply.github.com>
Date: Mon, 4 Nov 2024 10:23:39 -0500
Subject: [PATCH] feat(api-graphql): return entire payload in callback (#13980)

feat(api-graphql): return entire payload in callback
---
 .../__tests__/AWSAppSyncEventProvider.test.ts | 34 +++++++++++++++++++
 .../AWSAppSyncEventsProvider/index.ts         |  2 +-
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/packages/api-graphql/__tests__/AWSAppSyncEventProvider.test.ts b/packages/api-graphql/__tests__/AWSAppSyncEventProvider.test.ts
index d84f1daad89..fbb865be862 100644
--- a/packages/api-graphql/__tests__/AWSAppSyncEventProvider.test.ts
+++ b/packages/api-graphql/__tests__/AWSAppSyncEventProvider.test.ts
@@ -171,6 +171,40 @@ describe('AppSyncEventProvider', () => {
 						'Connection failed: Retriable Test',
 					);
 				});
+
+				test('subscription observer is triggered when a connection is formed and a data message is received after connection ack', async () => {
+					expect.assertions(1);
+					const mockNext = jest.fn();
+
+					const observer = provider.subscribe({
+						appSyncGraphqlEndpoint: 'ws://localhost:8080',
+					});
+
+					const event = JSON.stringify({ some: 'data' });
+
+					observer.subscribe({
+						// Succeed only when the first message comes through
+						next: mockNext,
+						// Closing a hot connection (for cleanup) makes it blow up the test stack
+						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),
+					});
+				});
 			});
 		});
 	});
diff --git a/packages/api-graphql/src/Providers/AWSAppSyncEventsProvider/index.ts b/packages/api-graphql/src/Providers/AWSAppSyncEventsProvider/index.ts
index 5bfeddd89fa..eff48289acc 100644
--- a/packages/api-graphql/src/Providers/AWSAppSyncEventsProvider/index.ts
+++ b/packages/api-graphql/src/Providers/AWSAppSyncEventsProvider/index.ts
@@ -167,7 +167,7 @@ export class AWSAppSyncEventProvider extends AWSWebSocketProvider {
 		if (type === MESSAGE_TYPES.DATA && payload) {
 			const deserializedEvent = JSON.parse(payload);
 			if (observer) {
-				observer.next(deserializedEvent);
+				observer.next({ id, type, event: deserializedEvent });
 			} else {
 				this.logger.debug(`observer not found for id: ${id}`);
 			}