fix(auth): issue identitied running integ tests

aws-amplify · Dec 28, 2023 · 51a0693 · 51a0693
1 parent 58b61f8
commit 51a0693
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 18 deletions.
diff --git a/packages/auth/__tests__/providers/cognito/utils/oauth/attemptCompleteOAuthFlow.test.ts b/packages/auth/__tests__/providers/cognito/utils/oauth/attemptCompleteOAuthFlow.test.ts
@@ -10,7 +10,6 @@ import { attemptCompleteOAuthFlow } from '../../../../../src/providers/cognito/u
 import { completeOAuthFlow } from '../../../../../src/providers/cognito/utils/oauth/completeOAuthFlow';
 import { getRedirectUrl } from '../../../../../src/providers/cognito/utils/oauth/getRedirectUrl';
 import { oAuthStore } from '../../../../../src/providers/cognito/utils/oauth/oAuthStore';
-import { addInflightPromise } from '../../../../../src/providers/cognito/utils/oauth/inflightPromise';
 import { cognitoUserPoolsTokenProvider } from '../../../../../src/providers/cognito/tokenProvider/tokenProvider';
 import { mockAuthConfigWithOAuth } from '../../../../mockData';
 
@@ -58,7 +57,6 @@ const mockAssertOAuthConfig = assertOAuthConfig as jest.Mock;
 const mockAssertTokenProviderConfig = assertTokenProviderConfig as jest.Mock;
 const mockCompleteOAuthFlow = completeOAuthFlow as jest.Mock;
 const mockGetRedirectUrl = getRedirectUrl as jest.Mock;
-const mockAddInflightPromise = addInflightPromise as jest.Mock;
 
 describe('attemptCompleteOAuthFlow', () => {
 	let windowSpy = jest.spyOn(window, 'window', 'get');

diff --git a/packages/auth/__tests__/providers/cognito/utils/oauth/completeOAuthFlow.test.ts b/packages/auth/__tests__/providers/cognito/utils/oauth/completeOAuthFlow.test.ts
@@ -4,14 +4,13 @@
 import { Hub, decodeJWT } from '@aws-amplify/core';
 import { handleFailure } from '../../../../../src/providers/cognito/utils/oauth/handleFailure';
 import { validateState } from '../../../../../src/providers/cognito/utils/oauth/validateState';
-// import { createOAuthError } from '../../../../../src/providers/cognito/utils/oauth/createOAuthError';
 import { resolveAndClearInflightPromises } from '../../../../../src/providers/cognito/utils/oauth/inflightPromise';
 import { oAuthStore } from '../../../../../src/providers/cognito/utils/oauth/oAuthStore';
 import { cacheCognitoTokens } from '../../../../../src/providers/cognito/tokenProvider/cacheTokens';
 import { AuthError } from '../../../../../src/errors/AuthError';
 import { AuthErrorTypes } from '../../../../../src/types/Auth';
 import { OAuthStore } from '../../../../../src/providers/cognito/utils/types';
-import { getCurrentUser } from '../../../../../src/providers/cognito/apis/getCurrentUser';
+import { cognitoUserPoolsTokenProvider } from '../../../../../src/providers/cognito/tokenProvider/tokenProvider';
 
 import { completeOAuthFlow } from '../../../../../src/providers/cognito/utils/oauth/completeOAuthFlow';
 
@@ -25,7 +24,6 @@ jest.mock('@aws-amplify/core', () => ({
 jest.mock('../../../../../src/providers/cognito/utils/oauth//handleFailure');
 jest.mock('../../../../../src/providers/cognito/utils/oauth/validateState');
 jest.mock('../../../../../src/providers/cognito/utils/oauth/inflightPromise');
-// jest.mock('../../../../../src/providers/cognito/utils/oauth/createOAuthError');
 jest.mock('../../../../../src/providers/cognito/apis/getCurrentUser');
 jest.mock('../../../../../src/providers/cognito/tokenProvider/cacheTokens');
 jest.mock(
@@ -46,6 +44,14 @@ jest.mock(
 		} as OAuthStore,
 	})
 );
+jest.mock(
+	'../../../../../src/providers/cognito/tokenProvider/tokenProvider',
+	() => ({
+		cognitoUserPoolsTokenProvider: {
+			setWaitForInflightOAuth: jest.fn(),
+		},
+	})
+);
 
 const mockHandleFailure = handleFailure as jest.Mock;
 const mockValidateState = validateState as jest.Mock;
@@ -84,6 +90,9 @@ describe('completeOAuthFlow', () => {
 		(oAuthStore.clearOAuthInflightData as jest.Mock).mockClear();
 		(oAuthStore.clearOAuthData as jest.Mock).mockClear();
 		(oAuthStore.storeOAuthSignIn as jest.Mock).mockClear();
+		(
+			cognitoUserPoolsTokenProvider.setWaitForInflightOAuth as jest.Mock
+		).mockClear();
 	});
 
 	it('handles error presented in the redirect url', async () => {
@@ -118,7 +127,7 @@ describe('completeOAuthFlow', () => {
 					...testInput,
 					currentUrl: `http://localhost:3000?state=someState123`,
 				})
-			).rejects.toThrow('The inflight OAuth flow has been cancelled.');
+			).rejects.toThrow('User cancelled OAuth flow.');
 		});
 
 		it('throws when `state` is not presented in the redirect url', async () => {
@@ -127,7 +136,7 @@ describe('completeOAuthFlow', () => {
 					...testInput,
 					currentUrl: `http://localhost:3000?code=123`,
 				})
-			).rejects.toThrow('The inflight OAuth flow has been cancelled.');
+			).rejects.toThrow('User cancelled OAuth flow.');
 		});
 
 		it('handles error when `validateState` fails', async () => {
@@ -221,6 +230,16 @@ describe('completeOAuthFlow', () => {
 			domain: 'oauth.domain.com',
 		};
 
+		it('throws when error and error_description are presented in the redirect url', () => {
+			const expectedErrorMessage = 'invalid_scope';
+			expect(
+				completeOAuthFlow({
+					...testInput,
+					currentUrl: `http://localhost:3000#error_description=${expectedErrorMessage}&error=invalid_request`,
+				})
+			).rejects.toThrow(expectedErrorMessage);
+		});
+
 		it('throws when access_token is not presented in the redirect url', () => {
 			expect(
 				completeOAuthFlow({
@@ -269,17 +288,25 @@ describe('completeOAuthFlow', () => {
 				ExpiresIn: expectedExpiresIn,
 			});
 
+			expect(oAuthStore.clearOAuthData).toHaveBeenCalledTimes(1);
+			expect(oAuthStore.storeOAuthSignIn).toHaveBeenCalledWith(true, undefined);
+			expect(mockResolveAndClearInflightPromises).toHaveBeenCalledTimes(1);
+			expect(
+				cognitoUserPoolsTokenProvider.setWaitForInflightOAuth
+			).toHaveBeenCalledTimes(1);
+
+			const waitForInflightOAuth = (
+				cognitoUserPoolsTokenProvider.setWaitForInflightOAuth as jest.Mock
+			).mock.calls[0][0];
+			expect(typeof waitForInflightOAuth).toBe('function');
+			expect(waitForInflightOAuth()).resolves.toBeUndefined();
+
+			expect(mockHubDispatch).toHaveBeenCalledTimes(2);
 			expect(mockReplaceState).toHaveBeenCalledWith(
 				{},
 				'',
 				testInput.redirectUri
 			);
-
-			expect(oAuthStore.clearOAuthData).toHaveBeenCalledTimes(1);
-			expect(oAuthStore.storeOAuthSignIn).toHaveBeenCalledWith(true, undefined);
-
-			expect(mockHubDispatch).toHaveBeenCalledTimes(2);
-			expect(mockResolveAndClearInflightPromises).toHaveBeenCalledTimes(1);
 		});
 	});
 });
diff --git a/packages/auth/src/providers/cognito/utils/oauth/completeOAuthFlow.ts b/packages/auth/src/providers/cognito/utils/oauth/completeOAuthFlow.ts
@@ -14,6 +14,7 @@ import { resolveAndClearInflightPromises } from './inflightPromise';
 import { cacheCognitoTokens } from '../../tokenProvider/cacheTokens';
 import { getCurrentUser } from '../../apis/getCurrentUser';
 import { createOAuthError } from './createOAuthError';
+import { cognitoUserPoolsTokenProvider } from '../../tokenProvider';
 
 export const completeOAuthFlow = async ({
 	currentUrl,
@@ -84,7 +85,7 @@ const handleCodeFlow = async ({
 	// 1. clicking the back button of browser
 	// 2. closing the provider hosted UI page and coming back to the app
 	if (!code || !state) {
-		throw createOAuthError('The inflight OAuth flow has been cancelled.');
+		throw createOAuthError('User cancelled OAuth flow.');
 	}
 
 	// may throw error is being caught in attemptCompleteOAuthFlow.ts
@@ -172,9 +173,8 @@ const handleImplicitFlow = async ({
 		state,
 		token_type,
 		expires_in,
-		// the following have been handled at line 40
-		// error_description,
-		// error,
+		error_description,
+		error,
 	} = (url.hash ?? '#')
 		.substring(1) // Remove # from returned code
 		.split('&')
@@ -189,6 +189,10 @@ const handleImplicitFlow = async ({
 			error: undefined,
 		});
 
+	if (error) {
+		throw createOAuthError(error_description ?? error);
+	}
+
 	if (!access_token) {
 		// error is being caught in attemptCompleteOAuthFlow.ts
 		throw createOAuthError('No access token returned from OAuth flow.');
@@ -224,6 +228,16 @@ const completeFlow = async ({
 }) => {
 	await oAuthStore.clearOAuthData();
 	await oAuthStore.storeOAuthSignIn(true, preferPrivateSession);
+
+	// this should be called before any call that involves `fetchAuthSession`
+	// e.g. `getCurrentUser()` below, so it allows every inflight async calls to
+	//  `fetchAuthSession` can be resolved
+	resolveAndClearInflightPromises();
+
+	// when the oauth flow is completed, there should be nothing to block the async calls
+	// that involves fetchAuthSession in the `TokenOrchestrator`
+	cognitoUserPoolsTokenProvider.setWaitForInflightOAuth(async () => {});
+
 	if (isCustomState(state)) {
 		Hub.dispatch(
 			'auth',
@@ -243,7 +257,6 @@ const completeFlow = async ({
 		AMPLIFY_SYMBOL
 	);
 	clearHistory(redirectUri);
-	resolveAndClearInflightPromises();
 };
 
 const isCustomState = (state: string): Boolean => {