updated graphql API singleton to use separate realtime providers per …

…endpoint
aws-amplify · Dec 18, 2024 · e61d28d · e61d28d
1 parent 7f6a612
commit e61d28d
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 11 deletions.
diff --git a/packages/api-graphql/__tests__/internals/generateClient.test.ts b/packages/api-graphql/__tests__/internals/generateClient.test.ts
@@ -270,7 +270,14 @@ describe('generateClient', () => {
 			const client = generateClient<Schema>({ amplify: Amplify });
 
 			const spy = jest.fn(() => from([graphqlMessage]));
-			(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
+			(raw.GraphQLAPI as any).appSyncRealTime = {
+				get() {
+					return { subscribe: spy }
+				},
+				set() {
+					// not needed for test mock
+				}
+			};
 
 			expect(normalizePostGraphqlCalls(spy)).toMatchSnapshot();
 
@@ -497,7 +504,14 @@ describe('generateClient', () => {
 				const client = generateClient<Schema>({ amplify: Amplify });
 
 				const spy = jest.fn(() => from([graphqlMessage]));
-				(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
+				(raw.GraphQLAPI as any).appSyncRealTime = {
+					get() {
+						return { subscribe: spy }
+					},
+					set() {
+						// not needed for test mock
+					}
+				};
 
 				client.models.Note.onCreate({
 					filter: graphqlVariables.filter,
@@ -531,7 +545,14 @@ describe('generateClient', () => {
 				});
 
 				const spy = jest.fn(() => from([graphqlMessage]));
-				(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
+				(raw.GraphQLAPI as any).appSyncRealTime = {
+					get() {
+						return { subscribe: spy }
+					},
+					set() {
+						// not needed for test mock
+					}
+				};
 
 				client.models.Note.onCreate({
 					filter: graphqlVariables.filter,
@@ -561,7 +582,14 @@ describe('generateClient', () => {
 				const client = generateClient<Schema>({ amplify: Amplify });
 
 				const spy = jest.fn(() => from([graphqlMessage]));
-				(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
+				(raw.GraphQLAPI as any).appSyncRealTime = {
+					get() {
+						return { subscribe: spy }
+					},
+					set() {
+						// not needed for test mock
+					}
+				};
 
 				client.models.Note.onCreate({
 					filter: graphqlVariables.filter,
@@ -583,7 +611,14 @@ describe('generateClient', () => {
 				const client = generateClient<Schema>({ amplify: Amplify });
 
 				const spy = jest.fn(() => from([graphqlMessage]));
-				(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
+				(raw.GraphQLAPI as any).appSyncRealTime = {
+					get() {
+						return { subscribe: spy }
+					},
+					set() {
+						// not needed for test mock
+					}
+				};
 
 				client.models.Note.onCreate({
 					filter: graphqlVariables.filter,
@@ -711,7 +746,14 @@ describe('generateClient', () => {
 			const client = generateClient<Schema>({ amplify: Amplify });
 
 			const spy = jest.fn(() => from([graphqlMessage]));
-			(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
+			(raw.GraphQLAPI as any).appSyncRealTime = {
+				get() {
+					return { subscribe: spy }
+				},
+				set() {
+					// not needed for test mock
+				}
+			};
 
 			client.models.Note.onCreate({
 				filter: graphqlVariables.filter,

diff --git a/packages/api-graphql/src/internals/InternalGraphQLAPI.ts b/packages/api-graphql/src/internals/InternalGraphQLAPI.ts
@@ -52,7 +52,7 @@ export class InternalGraphQLAPIClass {
 	/**
 	 * @private
 	 */
-	private appSyncRealTime = new AWSAppSyncRealTimeProvider();
+	private appSyncRealTime = new Map<string, AWSAppSyncRealTimeProvider>();
 
 	private _api = {
 		post,
@@ -384,12 +384,26 @@ export class InternalGraphQLAPIClass {
 		 */
 		const { headers: libraryConfigHeaders } = resolveLibraryOptions(amplify);
 
-		return this.appSyncRealTime
+		const appSyncGraphqlEndpoint = endpoint ?? config?.endpoint;
+
+		if (!appSyncGraphqlEndpoint) {
+			throw new Error(
+				'Endpoint missing from subscription query. An endpoint must either be configured or provided at the call site.',
+			);
+		}
+
+		const realtimeProvider =
+			this.appSyncRealTime.get(appSyncGraphqlEndpoint!) ??
+			new AWSAppSyncRealTimeProvider();
+
+		this.appSyncRealTime.set(appSyncGraphqlEndpoint, realtimeProvider);
+
+		return realtimeProvider
 			.subscribe(
 				{
 					query: print(query as DocumentNode),
 					variables,
-					appSyncGraphqlEndpoint: endpoint ?? config?.endpoint,
+					appSyncGraphqlEndpoint,
 					region: config?.region,
 					authenticationType: authMode,
 					apiKey: apiKeyOverride ?? config?.apiKey,

diff --git a/packages/api/__tests__/API.test.ts b/packages/api/__tests__/API.test.ts
@@ -33,7 +33,20 @@ const DEFAULT_AUTH_TOKEN =
 	'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyLCJleHAiOjE3MTAyOTMxMzB9.YzDpgJsrB3z-ZU1XxMcXSQsMbgCzwH_e-_76rnfehh0';
 
 const _postSpy = jest.spyOn((GraphQLAPI as any)._api, 'post');
-const _subspy = jest.spyOn((GraphQLAPI as any).appSyncRealTime, 'subscribe');
+const _subspy = jest.fn();
+
+/**
+ * Should be called on every subscription, ensuring that realtime provider instances
+ * are re-used for each distinct endpoint.
+ */
+const _setProviderSpy = jest.fn();
+
+(GraphQLAPI as any).appSyncRealTime = {
+	get() {
+		return { subscribe: _subspy }
+	},
+	set: _setProviderSpy
+};
 
 /**
  * Validates that a specific "post" occurred (against `_postSpy`).
@@ -116,6 +129,7 @@ function expectSubscription({
 		}),
 		expect.anything(),
 	);
+	expect(_setProviderSpy).toHaveBeenCalledWith(endpoint, expect.anything());
 }
 
 /**
@@ -204,7 +218,7 @@ describe('generateClient (web)', () => {
 	});
 
 	afterEach(() => {
-		jest.resetAllMocks();
+		jest.clearAllMocks();
 	});
 
 	for (const op of ['query', 'subscription'] as const) {