Datastore/feat user agent suffix (aws-amplify#10086)

* Adds suffix to user agent for calls to API initiated by DataStore

* Attempts to fix first half of user agent not being sent

* Makes setting of user agent header more concise

* Moves appending of suffix to user agent to core library

* Moves user agent suffix constant to common util in datastore

* Removes unused import

* Unit test for api-graphql

* Pulls in user agent suffix from datastore utils class

* Adds unit test for getAmplifyUserAgent with and without content

* Fixes issue found while testing, line too long.

* Adds test for DataStore mutation.ts

* Tests user agent suffix in datastore sync

* Adds user agent suffix assertion to subscription test

* Fixes variable declaration: const instead of let

* Removes leftover lines of code from testing objectContains

* Removes code style changes unrelated to user agent suffix

* Removes code style changes unrelated to user agent suffix

* Removes unnecessary null value option for userAgentSuffix - undefined is sufficient

* Replaces imports from lib-esm

* Replaces hard-coded string in assertion with constant from util file

* Moves var declaration under import

* Makes test method names more descriptive

Co-authored-by: Erin Beal <erinleig@amazon.com>

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

@@ -1237,7 +1237,63 @@ describe('API test', () => {
 				},
 			});
 		});
-	});
+
+		test('sends userAgent with suffix in request', async () => {
+			const spyonAuth = jest
+				.spyOn(Credentials, 'get')
+				.mockImplementationOnce(() => {
+					return new Promise((res, rej) => {
+						res('cred');
+					});
+				});
+
+			const spyon = jest
+				.spyOn(RestClient.prototype, 'post')
+				.mockImplementationOnce((url, init) => {
+					return new Promise((res, rej) => {
+						res({});
+					});
+				});
+
+			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' },
+				userAgentSuffix = '/DataStore';
+			api.configure({
+				aws_appsync_graphqlEndpoint: url,
+				aws_appsync_region: region,
+				aws_appsync_authenticationType: 'API_KEY',
+				aws_appsync_apiKey: apiKey,
+			});
+
+			const headers = {
+				Authorization: null,
+				'X-Api-Key': apiKey,
+				'x-amz-user-agent': `${Constants.userAgent}${userAgentSuffix}`,
+			};
+
+			const body = {
+				query: getEventQuery,
+				variables,
+			};
+
+			const init = {
+				headers,
+				body,
+				signerServiceInfo: {
+					service: 'appsync',
+					region,
+				},
+				cancellableToken: mockCancellableToken,
+			};
+			let authToken: undefined;
+
+			await api.graphql(graphqlOperation(GetEvent, variables, authToken, userAgentSuffix));
+
+			expect(spyon).toBeCalledWith(url, init);
+		});
 
 	describe('configure test', () => {
 		test('without aws_project_region', () => {

packages/api-graphql/src/GraphQLAPI.ts

@@ -22,8 +22,8 @@ import Observable from 'zen-observable-ts';
 import {
 	Amplify,
 	ConsoleLogger as Logger,
-	Constants,
 	Credentials,
+	getAmplifyUserAgent,
 	INTERNAL_AWS_APPSYNC_REALTIME_PUBSUB_PROVIDER,
 } from '@aws-amplify/core';
 import PubSub from '@aws-amplify/pubsub';
@@ -43,11 +43,13 @@ const logger = new Logger('GraphQLAPI');
 export const graphqlOperation = (
 	query,
 	variables = {},
-	authToken?: string
+	authToken?: string,
+	userAgentSuffix?: string
 ) => ({
 	query,
 	variables,
 	authToken,
+	userAgentSuffix,
 });
 
 /**
@@ -224,7 +226,13 @@ export class GraphQLAPIClass {
 	 * @returns An Observable if the query is a subscription query, else a promise of the graphql result.
 	 */
 	graphql<T = any>(
-		{ query: paramQuery, variables = {}, authMode, authToken }: GraphQLOptions,
+		{
+			query: paramQuery,
+			variables = {},
+			authMode,
+			authToken,
+			userAgentSuffix,
+		}: GraphQLOptions,
 		additionalHeaders?: { [key: string]: string }
 	): Observable<GraphQLResult<T>> | Promise<GraphQLResult<T>> {
 		const query =
@@ -233,7 +241,7 @@ export class GraphQLAPIClass {
 				: parse(print(paramQuery));
 
 		const [operationDef = {}] = query.definitions.filter(
-			(def) => def.kind === 'OperationDefinition'
+			def => def.kind === 'OperationDefinition'
 		);
 		const { operation: operationType } =
 			operationDef as OperationDefinitionNode;
@@ -251,7 +259,7 @@ export class GraphQLAPIClass {
 				const cancellableToken = this._api.getCancellableToken();
 				const initParams = { cancellableToken };
 				const responsePromise = this._graphql<T>(
-					{ query, variables, authMode },
+					{ query, variables, authMode, userAgentSuffix },
 					headers,
 					initParams
 				);
@@ -268,7 +276,7 @@ export class GraphQLAPIClass {
 	}
 
 	private async _graphql<T = any>(
-		{ query, variables, authMode }: GraphQLOptions,
+		{ query, variables, authMode, userAgentSuffix }: GraphQLOptions,
 		additionalHeaders = {},
 		initParams = {}
 	): Promise<GraphQLResult<T>> {
@@ -294,7 +302,7 @@ export class GraphQLAPIClass {
 			...(await graphql_headers({ query, variables })),
 			...additionalHeaders,
 			...(!customGraphqlEndpoint && {
-				[USER_AGENT_HEADER]: Constants.userAgent,
+				[USER_AGENT_HEADER]: getAmplifyUserAgent(userAgentSuffix),
 			}),
 		};
 
@@ -421,14 +429,14 @@ export class GraphQLAPIClass {
 	 */
 	_ensureCredentials() {
 		return this.Credentials.get()
-			.then((credentials) => {
+			.then(credentials => {
 				if (!credentials) return false;
 				const cred = this.Credentials.shear(credentials);
 				logger.debug('set credentials for api', cred);
 
 				return true;
 			})
-			.catch((err) => {
+			.catch(err => {
 				logger.warn('ensure credentials error', err);
 				return false;
 			});

packages/api-graphql/src/types/index.ts

@@ -20,6 +20,7 @@ export interface GraphQLOptions {
 	variables?: object;
 	authMode?: keyof typeof GRAPHQL_AUTH_MODE;
 	authToken?: string;
+	userAgentSuffix?: string;
 }
 
 export interface GraphQLResult<T = object> {

packages/core/__tests__/Platform-test.ts

@@ -1,3 +1,4 @@
+import { getAmplifyUserAgent } from '../src/Platform';
 import Platform from '../src/Platform';
 
 describe('Platform test', () => {
@@ -6,4 +7,16 @@ describe('Platform test', () => {
 			expect(Platform.isReactNative).toBe(false);
 		});
 	});
+
+	describe('getAmplifyUserAgent test', () => {
+		test('without content', () => {
+			expect(getAmplifyUserAgent()).toBe(Platform.userAgent);
+		});
+
+		test('with content', () => {
+			expect(getAmplifyUserAgent('/DataStore')).toBe(
+				`${Platform.userAgent}/DataStore`
+			);
+		});
+	});
 });

packages/core/src/Platform/index.ts

@@ -35,8 +35,8 @@ if (typeof navigator !== 'undefined' && navigator.product) {
 	}
 }
 
-export const getAmplifyUserAgent = () => {
-	return Platform.userAgent;
+export const getAmplifyUserAgent = (content?: string) => {
+	return `${Platform.userAgent}${content ? content : ''}`;
 };
 
 /**

packages/datastore/__tests__/mutation.test.ts

@@ -1,3 +1,4 @@
+const mockRestPost = jest.fn();
 import { RestClient } from '@aws-amplify/api-rest';
 import {
 	MutationProcessor,
@@ -17,6 +18,8 @@ import {
 } from '../src/types';
 import { createMutationInstanceFromModelOperation } from '../src/sync/utils';
 import { MutationEvent } from '../src/sync/';
+import { Constants } from '@aws-amplify/core';
+import { USER_AGENT_SUFFIX_DATASTORE } from '../src/util';
 
 let syncClasses: any;
 let modelInstanceCreator: any;
@@ -25,7 +28,11 @@ let PostCustomPK: PersistentModelConstructor<PostCustomPKType>;
 let PostCustomPKSort: PersistentModelConstructor<PostCustomPKSortType>;
 let axiosError;
 
-describe('Jittered retry', () => {
+beforeEach(() => {
+	axiosError = timeoutError;
+});
+
+describe('Jittered backoff', () => {
 	it('should progress exponentially until some limit', () => {
 		const COUNT = 13;
 
@@ -143,6 +150,21 @@ describe('MutationProcessor', () => {
 			expect(input.postId).toEqual(100);
 		});
 	});
+	describe('Call to rest api', () => {
+		it('Should send a user agent with the datastore suffix the rest api request', async () => {
+			jest.spyOn(mutationProcessor, 'resume');
+			await mutationProcessor.resume();
+
+			expect(mockRestPost).toBeCalledWith(
+				expect.anything(),
+				expect.objectContaining({
+					headers: expect.objectContaining({
+						'x-amz-user-agent': `${Constants.userAgent}${USER_AGENT_SUFFIX_DATASTORE}`,
+					}),
+				})
+			);
+		});
+	});
 	afterAll(() => {
 		jest.restoreAllMocks();
 	});
@@ -225,15 +247,14 @@ describe('error handler', () => {
 		);
 	});
 });
-
 // Mocking restClient.post to throw the error we expect
 // when experiencing poor network conditions
 jest.mock('@aws-amplify/api-rest', () => {
 	return {
 		...jest.requireActual('@aws-amplify/api-rest'),
 		RestClient() {
 			return {
-				post: jest.fn().mockImplementation(() => {
+				post: mockRestPost.mockImplementation(() => {
 					return Promise.reject(axiosError);
 				}),
 				getCancellableToken: () => {},
@@ -429,7 +450,3 @@ const timeoutError = {
 	},
 	code: 'ECONNABORTED',
 };
-
-beforeEach(() => {
-	axiosError = timeoutError;
-});

packages/datastore/__tests__/subscription.test.ts

@@ -1,7 +1,10 @@
+import Observable from 'zen-observable-ts';
+let mockObservable = new Observable(() => {});
+const mockGraphQL = jest.fn(() => mockObservable);
+
 import Amplify from 'aws-amplify';
 import { GRAPHQL_AUTH_MODE } from '@aws-amplify/api';
 import { CONTROL_MSG as PUBSUB_CONTROL_MSG } from '@aws-amplify/pubsub';
-import Observable from 'zen-observable-ts';
 import {
 	SubscriptionProcessor,
 	USER_CREDENTIALS,
@@ -16,14 +19,13 @@ import {
 	InternalSchema,
 	PersistentModelConstructor,
 } from '../src/types';
-
-let mockObservable = new Observable(() => {});
+import { USER_AGENT_SUFFIX_DATASTORE } from '../src/util';
 
 // mock graphql to return a mockable observable
 jest.mock('@aws-amplify/api', () => {
 	return {
 		...jest.requireActual('@aws-amplify/api'),
-		graphql: jest.fn(() => mockObservable),
+		graphql: mockGraphQL,
 	};
 });
 
@@ -650,6 +652,12 @@ describe('error handler', () => {
 							)
 						)
 					);
+
+					expect(mockGraphQL).toHaveBeenCalledWith(
+						expect.objectContaining({
+							userAgentSuffix: USER_AGENT_SUFFIX_DATASTORE,
+						})
+					);
 				});
 
 				done();

packages/datastore/__tests__/sync.test.ts

@@ -1,7 +1,8 @@
 // These tests should be replaced once SyncEngine.partialDataFeatureFlagEnabled is removed.
-
 import { GRAPHQL_AUTH_MODE } from '@aws-amplify/api-graphql';
 import { defaultAuthStrategy } from '../src/authModeStrategies';
+import { USER_AGENT_SUFFIX_DATASTORE } from '../src/util';
+let mockGraphQl;
 
 const sessionStorageMock = (() => {
 	let store = {};
@@ -281,6 +282,43 @@ describe('Sync', () => {
 				}
 			});
 		});
+
+		it('should send user agent suffix with graphql request', async () => {
+			window.sessionStorage.setItem('datastorePartialData', 'true');
+			const resolveResponse = {
+				data: {
+					syncPosts: {
+						items: [
+							{
+								id: '1',
+								title: 'Item 1',
+							},
+							{
+								id: '2',
+								title: 'Item 2',
+							},
+						],
+					},
+				},
+			};
+
+			const SyncProcessor = jitteredRetrySyncProcessorSetup({
+				resolveResponse,
+			});
+
+			await SyncProcessor.jitteredRetry({
+				query: defaultQuery,
+				variables: defaultVariables,
+				opName: defaultOpName,
+				modelDefinition: defaultModelDefinition,
+			});
+
+			expect(mockGraphQl).toHaveBeenCalledWith(
+				expect.objectContaining({
+					userAgentSuffix: USER_AGENT_SUFFIX_DATASTORE,
+				})
+			);
+		});
 	});
 
 	describe('error handler', () => {
@@ -409,16 +447,19 @@ function jitteredRetrySyncProcessorSetup({
 	coreMocks?: object;
 	errorHandler?: () => null;
 }) {
-	jest.mock('@aws-amplify/api', () => ({
-		...jest.requireActual('@aws-amplify/api'),
-		graphql: () =>
+	mockGraphQl = jest.fn(
+		() =>
 			new Promise((res, rej) => {
 				if (resolveResponse) {
 					res(resolveResponse);
 				} else if (rejectResponse) {
 					rej(rejectResponse);
 				}
-			}),
+			})
+	);
+	jest.mock('@aws-amplify/api', () => ({
+		...jest.requireActual('@aws-amplify/api'),
+		graphql: mockGraphQl,
 	}));
 
 	jest.mock('@aws-amplify/core', () => ({