custom header subscription tests

aws-amplify · Nov 11, 2023 · de5a0ca · de5a0ca
1 parent 5c029b2
commit de5a0ca
Show file tree

Hide file tree

Showing 6 changed files with 196 additions and 7 deletions.
diff --git a/.vscode/launch.json b/.vscode/launch.json
@@ -15,7 +15,7 @@
 				"--inspect-brk",
 				"${workspaceRoot}/node_modules/.bin/jest",
 				// Optionally specify a single test file to run/debug:
-				"GraphQLAPI.test.ts",
+				"generateClient.test.ts",
 				"--runInBand",
 				"--testTimeout",
 				"600000", // 10 min timeout so jest doesn't error while we're stepping through code

diff --git a/packages/api-graphql/__tests__/generateClient.test.ts b/packages/api-graphql/__tests__/generateClient.test.ts
@@ -3692,6 +3692,61 @@ describe('generateClient', () => {
 			);
 		});
 
+		test('can create() - with custom client header functions', async () => {
+			const spy = mockApiResponse({
+				data: {
+					createTodo: {
+						__typename: 'Todo',
+						...serverManagedFields,
+						name: 'some name',
+						description: 'something something',
+					},
+				},
+			});
+
+			const client = generateClient<Schema>({
+				amplify: Amplify,
+				headers: () => ({
+					'client-header-function': 'should return this header',
+				}),
+			});
+
+			const { data } = await client.models.Todo.create({
+				name: 'some name',
+				description: 'something something',
+			});
+
+			expect(spy).toHaveBeenCalledWith(
+				expect.objectContaining({
+					options: expect.objectContaining({
+						headers: expect.objectContaining({
+							'X-Api-Key': 'FAKE-KEY',
+							'client-header-function': 'should return this header',
+						}),
+						body: {
+							query: expect.stringContaining('createTodo(input: $input)'),
+							variables: {
+								input: {
+									name: 'some name',
+									description: 'something something',
+								},
+							},
+						},
+					}),
+				})
+			);
+
+			expect(data).toEqual(
+				expect.objectContaining({
+					__typename: 'Todo',
+					id: 'some-id',
+					owner: 'wirejobviously',
+					name: 'some name',
+					description: 'something something',
+				})
+			);
+		});
+
 		test('can create() - with custom request headers', async () => {
 			const spy = mockApiResponse({
 				data: {
@@ -3765,6 +3820,79 @@ describe('generateClient', () => {
 			);
 		});
 
+		test('can create() - with custom request header function', async () => {
+			const spy = mockApiResponse({
+				data: {
+					createTodo: {
+						__typename: 'Todo',
+						...serverManagedFields,
+						name: 'some name',
+						description: 'something something',
+					},
+				},
+			});
+
+			const client = generateClient<Schema>({
+				amplify: Amplify,
+				headers: {
+					'client-header': 'should not exist',
+				},
+			});
+
+			const { data } = await client.models.Todo.create(
+				{
+					name: 'some name',
+					description: 'something something',
+				},
+				{
+					headers: () => ({
+						'request-header-function': 'should return this header',
+					}),
+				}
+			);
+
+			expect(spy).toHaveBeenCalledWith(
+				expect.objectContaining({
+					options: expect.objectContaining({
+						headers: expect.objectContaining({
+							'X-Api-Key': 'FAKE-KEY',
+							'request-header-function': 'should return this header',
+						}),
+						body: {
+							query: expect.stringContaining('createTodo(input: $input)'),
+							variables: {
+								input: {
+									name: 'some name',
+									description: 'something something',
+								},
+							},
+						},
+					}),
+				})
+			);
+
+			// Request headers should overwrite client headers:
+			expect(spy).toHaveBeenCalledWith(
+				expect.objectContaining({
+					options: expect.objectContaining({
+						headers: expect.not.objectContaining({
+							'client-header': 'should not exist',
+						}),
+					}),
+				})
+			);
+
+			expect(data).toEqual(
+				expect.objectContaining({
+					__typename: 'Todo',
+					id: 'some-id',
+					owner: 'wirejobviously',
+					name: 'some name',
+					description: 'something something',
+				})
+			);
+		});
+
 		test('can get() - with custom client headers', async () => {
 			const spy = mockApiResponse({
 				data: {
@@ -4331,6 +4459,57 @@ describe('generateClient', () => {
 			});
 		});
 
+		test('can subscribe to onCreate() - with a custom header function', done => {
+			const noteToSend = {
+				__typename: 'Note',
+				...serverManagedFields,
+				body: 'a very good note',
+			};
+
+			const graphqlMessage = {
+				data: {
+					onCreateNote: noteToSend,
+				},
+			};
+
+			const graphqlVariables = {
+				filter: {
+					body: { contains: 'good note' },
+				},
+			};
+
+			const customHeaders = {
+				'subscription-header-function': 'should-return-this-header',
+			};
+
+			const customHeadersFunction = () => customHeaders;
+
+			const client = generateClient<Schema>({ amplify: Amplify });
+
+			const spy = jest.fn(() => from([graphqlMessage]));
+			(raw.GraphQLAPI as any).appSyncRealTime = { subscribe: spy };
+
+			client.models.Note.onCreate({
+				filter: graphqlVariables.filter,
+				headers: customHeadersFunction,
+			}).subscribe({
+				next(value) {
+					expectSubWithHeaders(
+						spy,
+						'onCreateNote',
+						graphqlVariables,
+						customHeaders
+					);
+					expect(value).toEqual(expect.objectContaining(noteToSend));
+					done();
+				},
+				error(error) {
+					expect(error).toBeUndefined();
+					done('bad news!');
+				},
+			});
+		});
+
 		test('can subscribe to onUpdate()', done => {
 			const noteToSend = {
 				__typename: 'Note',

diff --git a/packages/api-graphql/__tests__/utils/expects.ts b/packages/api-graphql/__tests__/utils/expects.ts
@@ -1,3 +1,5 @@
+import { CustomHeaders } from '@aws-amplify/data-schema-types';
+
 /**
  * Performs an `expect()` on a jest spy with some basic nested argument checks
  * based on the given mutation `opName` and `item`.
@@ -125,7 +127,7 @@ export function expectSubWithHeaders(
 	spy: jest.SpyInstance<any, any>,
 	opName: string,
 	item: Record<string, any>,
-	headers?: Record<string, string>
+	headers?: CustomHeaders
 ) {
 	expect(spy).toHaveBeenCalledWith(
 		expect.objectContaining({

diff --git a/packages/api-graphql/package.json b/packages/api-graphql/package.json
@@ -81,7 +81,7 @@
 	"dependencies": {
 		"@aws-amplify/core": "6.0.1",
 		"@aws-amplify/api-rest": "4.0.1",
-		"@aws-amplify/data-schema-types": "^0.4.2",
+		"@aws-amplify/data-schema-types": "^0.5.0",
 		"@aws-sdk/types": "3.387.0",
 		"graphql": "15.8.0",
 		"tslib": "^2.5.0",

diff --git a/packages/api-graphql/src/internals/InternalGraphQLAPI.ts b/packages/api-graphql/src/internals/InternalGraphQLAPI.ts
@@ -247,6 +247,14 @@ export class InternalGraphQLAPIClass {
 		const { headers: customHeaders, withCredentials } =
 			resolveLibraryOptions(amplify);
 
+		/**
+		 * Client or request-specific custom headers that may or may not be
+		 * returned by a function:
+		 */
+		if (typeof additionalHeaders === 'function') {
+			additionalHeaders = await additionalHeaders();
+		}
+
 		// TODO: Figure what we need to do to remove `!`'s.
 		const headers = {
 			...(!customEndpoint &&

diff --git a/packages/api-graphql/src/types/index.ts b/packages/api-graphql/src/types/index.ts
@@ -1,7 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 import { AmplifyClassV6, ResourcesConfig } from '@aws-amplify/core';
-import { ModelTypes } from '@aws-amplify/data-schema-types';
+import { ModelTypes, CustomHeaders } from '@aws-amplify/data-schema-types';
 import { Source, DocumentNode, GraphQLError } from 'graphql';
 export { OperationTypeNode } from 'graphql';
 import { Observable } from 'rxjs';
@@ -374,7 +374,7 @@ export type V6Client<T extends Record<any, any> = never> = ExcludeNeverFields<{
 	[__amplify]: AmplifyClassV6;
 	[__authMode]?: GraphQLAuthMode;
 	[__authToken]?: string;
-	[__headers]?: Record<string, string>;
+	[__headers]?: CustomHeaders;
 	graphql: GraphQLMethod;
 	cancel: (promise: Promise<any>, message?: string) => boolean;
 	isCancelError: (error: any) => boolean;
@@ -386,7 +386,7 @@ export type V6ClientSSRRequest<T extends Record<any, any> = never> =
 		[__amplify]: AmplifyClassV6;
 		[__authMode]?: GraphQLAuthMode;
 		[__authToken]?: string;
-		[__headers]?: Record<string, string>;
+		[__headers]?: CustomHeaders;
 		graphql: GraphQLMethodSSR;
 		cancel: (promise: Promise<any>, message?: string) => boolean;
 		isCancelError: (error: any) => boolean;
@@ -398,7 +398,7 @@ export type V6ClientSSRCookies<T extends Record<any, any> = never> =
 		[__amplify]: AmplifyClassV6;
 		[__authMode]?: GraphQLAuthMode;
 		[__authToken]?: string;
-		[__headers]?: Record<string, string>;
+		[__headers]?: CustomHeaders;
 		graphql: GraphQLMethod;
 		cancel: (promise: Promise<any>, message?: string) => boolean;
 		isCancelError: (error: any) => boolean;