fix(@aws-amplify/api-graphql) Respect additionalHeaders in GraphQLAPI…

… subscription requests (aws-amplify#5668)

* pass additionalHeaders to GraphQLAPI subscriptions

* add tests for additional headers in graphql subscriptions

Co-authored-by: Sam Martinez <samlmar@amazon.com>

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

@@ -858,6 +858,62 @@ describe('API test', () => {
 			expect(observable).not.toBe(undefined);
 		});
 
+		test('happy case subscription with additionalHeaders', async done => {
+			jest
+				.spyOn(RestClient.prototype, 'post')
+				.mockImplementation(async (url, init) => ({
+					extensions: {
+						subscription: {
+							newSubscriptions: {},
+						},
+					},
+				}));
+
+			const api = new API(config);
+			const url = 'https://appsync.amazonaws.com',
+				region = 'us-east-2',
+				apiKey = 'secret_api_key',
+				variables = { id: '809392da-ec91-4ef0-b219-5238a8f942b2' };
+
+			api.configure({
+				aws_appsync_graphqlEndpoint: url,
+				aws_appsync_region: region,
+				aws_appsync_authenticationType: 'API_KEY',
+				aws_appsync_apiKey: apiKey,
+			});
+
+			PubSub.subscribe = jest.fn(() => Observable.of({}));
+
+			const SubscribeToEventComments = `subscription SubscribeToEventComments($eventId: String!) {
+				subscribeToEventComments(eventId: $eventId) {
+					eventId
+					commentId
+					content
+				}
+			}`;
+
+			const doc = parse(SubscribeToEventComments);
+			const query = print(doc);
+
+			const additionalHeaders = {
+				'x-custom-header': 'value',
+			};
+
+			const observable = (api.graphql(
+				graphqlOperation(query, variables),
+				additionalHeaders
+			) as Observable<object>).subscribe({
+				next: () => {
+					expect(PubSub.subscribe).toHaveBeenCalledTimes(1);
+					const subscribeOptions = (PubSub.subscribe as any).mock.calls[0][1];
+					expect(subscribeOptions.additionalHeaders).toBe(additionalHeaders);
+					done();
+				},
+			});
+
+			expect(observable).not.toBe(undefined);
+		});
+
 		test('happy case mutation', async () => {
 			const spyonAuth = jest
 				.spyOn(Credentials, 'get')

packages/api-graphql/src/GraphQLAPI.ts

@@ -198,7 +198,10 @@ export class GraphQLAPIClass {
 			case 'mutation':
 				return this._graphql({ query, variables, authMode }, additionalHeaders);
 			case 'subscription':
-				return this._graphqlSubscribe({ query, variables, authMode });
+				return this._graphqlSubscribe(
+					{ query, variables, authMode },
+					additionalHeaders
+				);
 		}
 
 		throw new Error(`invalid operation type: ${operationType}`);
@@ -277,11 +280,10 @@ export class GraphQLAPIClass {
 		return response;
 	}
 
-	private _graphqlSubscribe({
-		query,
-		variables,
-		authMode: defaultAuthenticationType,
-	}: GraphQLOptions): Observable<any> {
+	private _graphqlSubscribe(
+		{ query, variables, authMode: defaultAuthenticationType }: GraphQLOptions,
+		additionalHeaders = {}
+	): Observable<any> {
 		const {
 			aws_appsync_region: region,
 			aws_appsync_graphqlEndpoint: appSyncGraphqlEndpoint,
@@ -302,6 +304,7 @@ export class GraphQLAPIClass {
 				region,
 				variables,
 				graphql_headers,
+				additionalHeaders,
 			});
 		} else {
 			logger.debug('No pubsub module applied for subscription');

packages/pubsub/src/Providers/AWSAppSyncRealTimeProvider.ts

@@ -234,6 +234,7 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 			apiKey,
 			region,
 			graphql_headers = () => ({}),
+			additionalHeaders = {},
 		} = options;
 
 		const subscriptionState: SUBSCRIPTION_STATUS = SUBSCRIPTION_STATUS.PENDING;
@@ -263,6 +264,7 @@ export class AWSAppSyncRealTimeProvider extends AbstractPubSubProvider {
 				region,
 			})),
 			...(await graphql_headers()),
+			...additionalHeaders,
 			[USER_AGENT_HEADER]: Constants.userAgent,
 		};