fix(aws-amplify): destroyAmplifyServerContext may not be called

aws-amplify · Nov 22, 2023 · b1552f6 · b1552f6
1 parent b877493
commit b1552f6
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 5 deletions.
diff --git a/packages/aws-amplify/__tests__/adapterCore/runWithAmplifyServerContext.test.ts b/packages/aws-amplify/__tests__/adapterCore/runWithAmplifyServerContext.test.ts
@@ -29,6 +29,10 @@ describe('runWithAmplifyServerContext', () => {
 		mockCreateAmplifyServerContext.mockReturnValueOnce(mockContextSpec);
 	});
 
+	afterEach(() => {
+		mockDestroyAmplifyServerContext.mockReset();
+	});
+
 	it('should run the operation with the context', () => {
 		const mockOperation = jest.fn();
 		runWithAmplifyServerContext(
@@ -45,7 +49,7 @@ describe('runWithAmplifyServerContext', () => {
 		expect(mockOperation).toHaveBeenCalledWith(mockContextSpec);
 	});
 
-	it('should destroy the context after the operation', async () => {
+	it('should destroy the context after the operation completed', async () => {
 		const mockOperation = jest.fn();
 		await runWithAmplifyServerContext(
 			mockAmplifyConfig,
@@ -63,6 +67,29 @@ describe('runWithAmplifyServerContext', () => {
 		);
 	});
 
+	it('should destroy the context when the operation throws', async () => {
+		const testError = new Error('some error');
+		const mockOperation = jest.fn();
+		mockOperation.mockRejectedValueOnce(testError);
+
+		await expect(
+			runWithAmplifyServerContext(
+				mockAmplifyConfig,
+				{
+					Auth: {
+						tokenProvider: mockTokenProvider,
+						credentialsProvider: mockCredentialAndIdentityProvider,
+					},
+				},
+				mockOperation
+			)
+		).rejects.toThrow(testError);
+
+		expect(mockDestroyAmplifyServerContext).toHaveBeenCalledWith(
+			mockContextSpec
+		);
+	});
+
 	it('should return the result returned by the operation callback function', async () => {
 		const mockResultValue = {
 			url: 'http://123.com',

diff --git a/packages/aws-amplify/src/adapterCore/runWithAmplifyServerContext.ts b/packages/aws-amplify/src/adapterCore/runWithAmplifyServerContext.ts
@@ -26,9 +26,12 @@ export const runWithAmplifyServerContext: AmplifyServer.RunOperationWithContext
 		);
 
 		// run the operation with injecting the context
-		const result = await operation(contextSpec);
+		try {
+			const result = await operation(contextSpec);
 
-		destroyAmplifyServerContext(contextSpec);
-
-		return result;
+			return result;
+		} finally {
+			// ensures destroy the context regardless whether the operation succeeded or failed
+			destroyAmplifyServerContext(contextSpec);
+		}
 	};