feat(data): add publish events over WebSocket (#14242)

Author: zxl629

packages/api-graphql/__tests__/AWSAppSyncEventProvider.test.ts

@@ -9,6 +9,56 @@ import { ConnectionState as CS } from '../src/types/PubSub';
 
 import { AWSAppSyncEventProvider } from '../src/Providers/AWSAppSyncEventsProvider';
 
+// Mock all calls to signRequest
+jest.mock('@aws-amplify/core/internals/aws-client-utils', () => {
+	const original = jest.requireActual(
+		'@aws-amplify/core/internals/aws-client-utils',
+	);
+	return {
+		...original,
+		signRequest: (_request: any, _options: any) => {
+			return {
+				method: 'test',
+				headers: { test: 'test' },
+				url: new URL('http://example/'),
+			};
+		},
+	};
+});
+
+// Mock all calls to signRequest
+jest.mock('@aws-amplify/core', () => {
+	const original = jest.requireActual('@aws-amplify/core');
+	const session = {
+		tokens: {
+			accessToken: {
+				toString: () => 'test',
+			},
+		},
+		credentials: {
+			accessKeyId: 'test',
+			secretAccessKey: 'test',
+		},
+	};
+	return {
+		...original,
+		fetchAuthSession: (_request: any, _options: any) => {
+			return Promise.resolve(session);
+		},
+		Amplify: {
+			Auth: {
+				fetchAuthSession: async () => session,
+			},
+		},
+		browserOrNode() {
+			return {
+				isBrowser: true,
+				isNode: false,
+			};
+		},
+	};
+});
+
 describe('AppSyncEventProvider', () => {
 	describe('subscribe()', () => {
 		describe('returned observer', () => {
@@ -245,4 +295,210 @@ describe('AppSyncEventProvider', () => {
 			});
 		});
 	});
+	describe('publish', () => {
+		let fakeWebSocketInterface: FakeWebSocketInterface;
+		const loggerSpy: jest.SpyInstance = jest.spyOn(
+			ConsoleLogger.prototype,
+			'_log',
+		);
+		let provider: AWSAppSyncEventProvider;
+		let reachabilityObserver: Observer<{ online: boolean }>;
+
+		beforeEach(async () => {
+			// Set the network to "online" for these tests
+			jest
+				.spyOn(Reachability.prototype, 'networkMonitor')
+				.mockImplementationOnce(() => {
+					return new Observable(observer => {
+						reachabilityObserver = observer;
+					});
+				})
+				// Twice because we subscribe to get the initial state then again to monitor reachability
+				.mockImplementationOnce(() => {
+					return new Observable(observer => {
+						reachabilityObserver = observer;
+					});
+				});
+
+			fakeWebSocketInterface = new FakeWebSocketInterface();
+			provider = new AWSAppSyncEventProvider();
+			Object.defineProperty(provider, 'socketStatus', {
+				value: constants.SOCKET_STATUS.READY,
+			});
+
+			Object.defineProperty(provider, 'awsRealTimeSocket', {
+				get: function () {
+					fakeWebSocketInterface.newWebSocket();
+					return fakeWebSocketInterface.webSocket;
+				},
+			});
+
+			// Saving this spy and resetting it by hand causes badness
+			//     Saving it causes new websockets to be reachable across past tests that have not fully closed
+			//     Resetting it proactively causes those same past tests to be dealing with null while they reach a settled state
+
+			// Set the addEventListener() to do nothing for the test
+			jest
+				.spyOn(fakeWebSocketInterface.webSocket, 'addEventListener')
+				.mockImplementation(() => {});
+
+			jest.spyOn(provider as any, '_getNewWebSocket').mockImplementation(() => {
+				fakeWebSocketInterface.newWebSocket();
+				return fakeWebSocketInterface.webSocket as WebSocket;
+			});
+
+			// Reduce retry delay for tests to 100ms
+			Object.defineProperty(constants, 'MAX_DELAY_MS', {
+				value: 100,
+			});
+			// Reduce retry delay for tests to 100ms
+			Object.defineProperty(constants, 'RECONNECT_DELAY', {
+				value: 100,
+			});
+		});
+
+		afterEach(async () => {
+			provider?.close();
+			await fakeWebSocketInterface?.closeInterface();
+			fakeWebSocketInterface?.teardown();
+			loggerSpy.mockClear();
+		});
+
+		test('authenticating with AWS_IAM', async () => {
+			expect.assertions(1);
+			provider.publish({
+				appSyncGraphqlEndpoint: 'ws://localhost:8080',
+				query: 'events/test',
+				variables: { some: 'data' },
+				authenticationType: 'iam',
+			});
+
+			await fakeWebSocketInterface?.readyForUse;
+
+			expect(loggerSpy).toHaveBeenCalledWith(
+				'DEBUG',
+				'Authenticating with "iam"',
+			);
+		});
+
+		test('authenticating with API_KEY', async () => {
+			expect.assertions(1);
+			fakeWebSocketInterface.webSocket.readyState = WebSocket.OPEN;
+			provider.publish({
+				appSyncGraphqlEndpoint: 'ws://localhost:8080',
+				query: 'events/test',
+				variables: { some: 'data' },
+				authenticationType: 'apiKey',
+				apiKey: 'test',
+			});
+
+			await fakeWebSocketInterface?.readyForUse;
+
+			expect(loggerSpy).toHaveBeenCalledWith(
+				'DEBUG',
+				'Authenticating with "apiKey"',
+			);
+		});
+
+		test('authenticating with OPENID_CONNECT', async () => {
+			expect.assertions(1);
+
+			provider.publish({
+				appSyncGraphqlEndpoint: 'ws://localhost:8080',
+				query: 'events/test',
+				variables: { some: 'data' },
+				authenticationType: 'oidc',
+			});
+
+			await fakeWebSocketInterface?.readyForUse;
+
+			expect(loggerSpy).toHaveBeenCalledWith(
+				'DEBUG',
+				'Authenticating with "oidc"',
+			);
+		});
+
+		test('authenticating with AWS_LAMBDA/custom', async () => {
+			expect.assertions(1);
+
+			provider.publish({
+				appSyncGraphqlEndpoint: 'ws://localhost:8080',
+				query: 'events/test',
+				variables: { some: 'data' },
+				authenticationType: 'none',
+				additionalHeaders: {
+					Authorization: 'test',
+				},
+			});
+
+			await fakeWebSocketInterface?.readyForUse;
+
+			expect(loggerSpy).toHaveBeenCalledWith(
+				'DEBUG',
+				'Authenticating with "none"',
+			);
+		});
+
+		test('authenticating with userPool / custom library options token', async () => {
+			expect.assertions(1);
+
+			provider.publish({
+				appSyncGraphqlEndpoint: 'ws://localhost:8080',
+				query: 'events/test',
+				variables: { some: 'data' },
+				authenticationType: 'userPool',
+				/**
+				 * When Amplify is configured with a `header` function
+				 * that returns an `Authorization` token, the GraphQL
+				 * API will pass this function as the `libraryConfigHeaders`
+				 * option to the AWSAppSyncRealTimeProvider's `subscribe`
+				 * function.
+				 */
+				libraryConfigHeaders: async () => ({
+					Authorization: 'test',
+				}),
+			});
+
+			await fakeWebSocketInterface?.readyForUse;
+
+			expect(loggerSpy).toHaveBeenCalledWith(
+				'DEBUG',
+				'Authenticating with "userPool"',
+			);
+		});
+
+		test('authenticating with AWS_LAMBDA/custom w/ custom header function', async () => {
+			expect.assertions(1);
+
+			provider.publish({
+				appSyncGraphqlEndpoint: 'ws://localhost:8080',
+				query: 'events/test',
+				variables: { some: 'data' },
+				authenticationType: 'none',
+				additionalHeaders: async () => ({
+					Authorization: 'test',
+				}),
+			});
+
+			await fakeWebSocketInterface?.readyForUse;
+
+			expect(loggerSpy).toBeCalledWith('DEBUG', 'Authenticating with "none"');
+		});
+
+		test('authenticating with AWS_LAMBDA/custom without Authorization', async () => {
+			expect.assertions(1);
+
+			await expect(
+				provider.publish({
+					appSyncGraphqlEndpoint: 'ws://localhost:8080',
+					query: 'events/test',
+					variables: { some: 'data' },
+					authenticationType: 'none',
+					additionalHeaders: {
+						Authorization: '',
+					},
+				}),
+			).rejects.toThrow('No auth token specified');
+		});
+	});
 });

packages/api-graphql/__tests__/events.test.ts

@@ -173,6 +173,42 @@ describe('Events client', () => {
 			});
 		});
 
+		describe('publish', () => {
+			test('happy publish', async () => {
+				const channel = await events.connect('/');
+
+				channel.publish({ some: 'data' });
+			});
+
+			test('publish() becomes invalid after .close() is called', async () => {
+				const channel = await events.connect('/');
+				channel.close();
+				await expect(channel.publish({ some: 'data' })).rejects.toThrow(
+					'Channel is closed',
+				);
+			});
+
+			describe('auth modes', () => {
+				let mockProvider: typeof AppSyncEventProvider;
+
+				beforeEach(() => {
+					mockProvider = AppSyncEventProvider;
+				});
+
+				for (const authMode of authModes) {
+					test(`auth override: ${authMode}`, async () => {
+						const channel = await events.connect('/');
+
+						channel.publish({ some: 'data' }, { authMode });
+
+						expect(mockProvider.publish).toHaveBeenCalledWith(
+							expect.objectContaining({ authenticationType: authMode }),
+						);
+					});
+				}
+			});
+		});
+
 		describe('post', () => {
 			let mockReq: typeof appsyncRequest;
 

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

@@ -12,6 +12,7 @@ import { CustomHeaders } from '@aws-amplify/data-schema/runtime';
 import { DEFAULT_KEEP_ALIVE_TIMEOUT, MESSAGE_TYPES } from '../constants';
 import { AWSWebSocketProvider } from '../AWSWebSocketProvider';
 import { awsRealTimeHeaderBasedAuth } from '../AWSWebSocketProvider/authHeaders';
+import { serializeEvents } from '../../internals/events/utils';
 
 // resolved/actual AuthMode values. identityPool gets resolves to IAM upstream in InternalGraphQLAPI._graphqlSubscribe
 type ResolvedGraphQLAuthModes = Exclude<GraphQLAuthMode, 'identityPool'>;
@@ -52,6 +53,7 @@ export class AWSAppSyncEventProvider extends AWSWebSocketProvider {
 			wsProtocolName: WS_PROTOCOL_NAME,
 			connectUri: CONNECT_URI,
 		});
+		this.allowNoSubscriptions = true;
 	}
 
 	getProviderName() {
@@ -73,7 +75,11 @@ export class AWSAppSyncEventProvider extends AWSWebSocketProvider {
 		options: AWSAppSyncEventProviderOptions,
 		customUserAgentDetails?: CustomUserAgentDetails,
 	) {
-		super.publish(options, customUserAgentDetails);
+		return super.publish(options, customUserAgentDetails);
+	}
+
+	public closeIfNoActiveSubscription() {
+		this._closeSocketIfRequired();
 	}
 
 	protected async _prepareSubscriptionPayload({
@@ -97,14 +103,14 @@ export class AWSAppSyncEventProvider extends AWSWebSocketProvider {
 			query,
 			apiKey,
 			region,
+			variables,
 		} = options;
 
-		// This will be needed for WS publish
-		// const data = {
-		// 	events: [variables],
-		// };
-
-		const serializedData = JSON.stringify({ channel: query });
+		const data = {
+			channel: query,
+			events: variables !== undefined ? serializeEvents(variables) : undefined,
+		};
+		const serializedData = JSON.stringify(data);
 
 		const headers = {
 			...(await awsRealTimeHeaderBasedAuth({
@@ -121,22 +127,23 @@ export class AWSAppSyncEventProvider extends AWSWebSocketProvider {
 			[USER_AGENT_HEADER]: getAmplifyUserAgent(customUserAgentDetails),
 		};
 
-		// Commented out code will be needed for WS publish
 		const subscriptionMessage = {
 			id: subscriptionId,
 			channel: query,
-			// events: [JSON.stringify(variables)],
+			events: variables !== undefined ? serializeEvents(variables) : undefined,
 			authorization: {
 				...headers,
 			},
-			// payload: {
-			// 	events: serializedData,
-			// 	extensions: {
-			// 		authorization: {
-			// 			...headers,
-			// 		},
-			// 	},
-			// },
+			payload: {
+				events:
+					variables !== undefined ? serializeEvents(variables) : undefined,
+				channel: query,
+				extensions: {
+					authorization: {
+						...headers,
+					},
+				},
+			},
 			type: publish
 				? MESSAGE_TYPES.EVENT_PUBLISH
 				: MESSAGE_TYPES.EVENT_SUBSCRIBE,