refactor(api-graphql): set observers earlier to ensure handling unsubscription properly (#14290)

* refactor(api-graphql): set observers earlier to ensure handling unsubscription properly

refs: aws-amplify/amplify-data#554

* Update to keep the set call closer to its original location

---------

Co-authored-by: Aaron S <[email protected]>
Co-authored-by: Aaron S. <[email protected]>

Author: no6v

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

@@ -992,6 +992,28 @@ describe('AWSAppSyncRealTimeProvider', () => {
 					expect(fakeWebSocketInterface.hasClosed).resolves.toBeUndefined();
 				});
 
+				test('unsubscription message should be sent even if unsubscribe immediately', async () => {
+					expect.assertions(1);
+
+					const sendUnsubscriptionMessageSpy = jest
+						.spyOn(provider as any, '_sendUnsubscriptionMessage');
+
+					provider
+						.subscribe({
+							appSyncGraphqlEndpoint: 'ws://localhost:8080',
+						})
+						.subscribe({ error: () => {} })
+						.unsubscribe();
+
+					await fakeWebSocketInterface?.standardConnectionHandshake();
+					await fakeWebSocketInterface?.startAckMessage();
+
+					expect(sendUnsubscriptionMessageSpy).toHaveBeenNthCalledWith(
+						1,
+						fakeWebSocketInterface?.webSocket.subscriptionId,
+					);
+				});
+
 				test('failure to ack before timeout', async () => {
 					expect.assertions(1);
 

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

@@ -158,6 +158,7 @@ export abstract class AWSWebSocketProvider {
 
 			let subscriptionStartInProgress = false;
 			const subscriptionId = amplifyUuid();
+
 			const startSubscription = () => {
 				if (!subscriptionStartInProgress) {
 					subscriptionStartInProgress = true;
@@ -396,9 +397,6 @@ export abstract class AWSWebSocketProvider {
 	}) {
 		const { query, variables } = options;
 
-		const { additionalCustomHeaders, libraryConfigHeaders } =
-			await additionalHeadersFromOptions(options);
-
 		this.subscriptionObserverMap.set(subscriptionId, {
 			observer,
 			query: query ?? '',
@@ -407,6 +405,9 @@ export abstract class AWSWebSocketProvider {
 			startAckTimeoutId: undefined,
 		});
 
+		const { additionalCustomHeaders, libraryConfigHeaders } =
+			await additionalHeadersFromOptions(options);
+
 		const serializedSubscriptionMessage =
 			await this._prepareSubscriptionPayload({
 				options,