fix more tests

aws · mergify · Jul 6, 2023 · Jun 7, 2023 · Jun 9, 2023 · Jun 9, 2023
commit 50729653b734356979aa9e3b41a9e779b20b457b
diff --git a/packages/aws-cdk-lib/aws-dynamodb/test/replica-provider.test.ts b/packages/aws-cdk-lib/aws-dynamodb/test/replica-provider.test.ts
@@ -1,4 +1,5 @@
-import * as AWS from 'aws-sdk-mock';
+import { DescribeTableCommand, DynamoDBClient, UpdateTableCommand } from '@aws-sdk/client-dynamodb';
+import { mockClient } from 'aws-sdk-client-mock';
 import * as sinon from 'sinon';
 import { OnEventRequest } from '../../custom-resources/lib/provider-framework/types';
 import { isCompleteHandler, onEventHandler } from '../lib/replica-handler';
@@ -14,8 +15,6 @@ afterAll(() => {
   global.console.log = oldConsoleLog;
 });
 
-AWS.setSDK(require.resolve('aws-sdk'));
-
 const REGION = 'eu-west-2';
 
 const createEvent: OnEventRequest = {
@@ -33,14 +32,16 @@ const createEvent: OnEventRequest = {
   ResourceType: 'resource-type',
 };
 
+const ddbMock = mockClient(DynamoDBClient);
+
 afterEach(() => {
-  AWS.restore();
+  ddbMock.reset();
 });
 
 test('on event', async () => {
   const updateTableMock = sinon.fake.resolves({});
 
-  AWS.mock('DynamoDB', 'updateTable', updateTableMock);
+  ddbMock.on(UpdateTableCommand).callsFake(updateTableMock);
 
   const data = await onEventHandler(createEvent);
 
@@ -61,16 +62,16 @@ test('on event', async () => {
 });
 
 test("on Update event from CFN calls updateTable with Create if a replica in the region doesn't exist", async () => {
-  AWS.mock('DynamoDB', 'describeTable', sinon.fake.resolves({
+  ddbMock.on(DescribeTableCommand).resolves({
     Table: {
       Replicas: [
         // no replicas exist yet
       ],
     },
-  }));
+  });
 
   const updateTableMock = sinon.fake.resolves({});
-  AWS.mock('DynamoDB', 'updateTable', updateTableMock);
+  ddbMock.on(UpdateTableCommand).callsFake(updateTableMock);
 
   const data = await onEventHandler({
     ...createEvent,
@@ -97,16 +98,16 @@ test("on Update event from CFN calls updateTable with Create if a replica in the
 });
 
 test("on Update event from CFN calls doesn't call updateTable if a replica in the region does exist", async () => {
-  AWS.mock('DynamoDB', 'describeTable', sinon.fake.resolves({
+  ddbMock.on(DescribeTableCommand).resolves({
     Table: {
       Replicas: [
         { RegionName: REGION },
       ],
     },
-  }));
+  });
 
   const updateTableMock = sinon.fake.resolves({});
-  AWS.mock('DynamoDB', 'updateTable', updateTableMock);
+  ddbMock.on(UpdateTableCommand).callsFake(updateTableMock);
 
   await onEventHandler({
     ...createEvent,
@@ -122,7 +123,7 @@ test("on Update event from CFN calls doesn't call updateTable if a replica in th
 test('on event calls updateTable with Delete', async () => {
   const updateTableMock = sinon.fake.resolves({});
 
-  AWS.mock('DynamoDB', 'updateTable', updateTableMock);
+  ddbMock.on(UpdateTableCommand).callsFake(updateTableMock);
 
   const data = await onEventHandler({
     ...createEvent,
@@ -149,7 +150,7 @@ test('is complete for create returns false without replicas', async () => {
     Table: {},
   });
 
-  AWS.mock('DynamoDB', 'describeTable', describeTableMock);
+  ddbMock.on(DescribeTableCommand).callsFake(describeTableMock);
 
   const data = await isCompleteHandler(createEvent);
 
@@ -168,7 +169,7 @@ test('is complete for create returns false when replica is not active', async ()
     },
   });
 
-  AWS.mock('DynamoDB', 'describeTable', describeTableMock);
+  ddbMock.on(DescribeTableCommand).callsFake(describeTableMock);
 
   const data = await isCompleteHandler(createEvent);
 
@@ -188,7 +189,7 @@ test('is complete for create returns false when table is not active', async () =
     },
   });
 
-  AWS.mock('DynamoDB', 'describeTable', describeTableMock);
+  ddbMock.on(DescribeTableCommand).callsFake(describeTableMock);
 
   const data = await isCompleteHandler(createEvent);
 
@@ -208,7 +209,7 @@ test('is complete for create returns true when replica is active', async () => {
     },
   });
 
-  AWS.mock('DynamoDB', 'describeTable', describeTableMock);
+  ddbMock.on(DescribeTableCommand).callsFake(describeTableMock);
 
   const data = await isCompleteHandler(createEvent);
 
@@ -228,7 +229,7 @@ test('is complete for delete returns true when replica is gone', async () => {
     },
   });
 
-  AWS.mock('DynamoDB', 'describeTable', describeTableMock);
+  ddbMock.on(DescribeTableCommand).callsFake(describeTableMock);
 
   const data = await isCompleteHandler({
     ...createEvent,

diff --git a/packages/aws-cdk-lib/aws-route53/test/cross-account-zone-delegation-handler/index.test.ts b/packages/aws-cdk-lib/aws-route53/test/cross-account-zone-delegation-handler/index.test.ts
@@ -17,15 +17,15 @@ jest.mock('@aws-sdk/client-sts', () => {
   return {
     STS: jest.fn().mockImplementation(() => {
       return mockStsClient;;
-    })
+    }),
   };
 });
 
 jest.mock('@aws-sdk/client-route-53', () => {
   return {
     Route53: jest.fn().mockImplementation(() => {
       return mockRoute53Client;
-    })
+    }),
   };
 });
 

diff --git a/packages/aws-cdk-lib/aws-route53/test/delete-existing-record-set-handler.test.ts b/packages/aws-cdk-lib/aws-route53/test/delete-existing-record-set-handler.test.ts
@@ -1,26 +1,21 @@
+import { handler } from '../lib/delete-existing-record-set-handler';
+
 const mockListResourceRecordSetsResponse = jest.fn();
 const mockChangeResourceRecordSetsResponse = jest.fn();
+const mockWaitUntilResourceRecordSetsChanged = jest.fn().mockResolvedValue({});
 
 const mockRoute53 = {
-  listResourceRecordSets: jest.fn().mockReturnValue({
-    promise: mockListResourceRecordSetsResponse,
-  }),
-  changeResourceRecordSets: jest.fn().mockReturnValue({
-    promise: mockChangeResourceRecordSetsResponse,
-  }),
-  waitFor: jest.fn().mockReturnValue({
-    promise: jest.fn().mockResolvedValue({}),
-  }),
+  listResourceRecordSets: mockListResourceRecordSetsResponse,
+  changeResourceRecordSets: mockChangeResourceRecordSetsResponse,
 };
 
-jest.mock('aws-sdk', () => {
+jest.mock('@aws-sdk/client-route-53', () => {
   return {
-    Route53: jest.fn(() => mockRoute53),
+    Route53: jest.fn().mockImplementation(() => mockRoute53),
+    waitUntilResourceRecordSetsChanged: (...args: any[]) => mockWaitUntilResourceRecordSetsChanged(...args),
   };
 });
 
-import { handler } from '../lib/delete-existing-record-set-handler';
-
 const event: AWSLambda.CloudFormationCustomResourceEvent & { PhysicalResourceId?: string } = {
   RequestType: 'Create',
   ServiceToken: 'service-token',
@@ -87,7 +82,8 @@ test('create request with existing record', async () => {
     },
   });
 
-  expect(mockRoute53.waitFor).toHaveBeenCalledWith('resourceRecordSetsChanged', {
+  expect(mockWaitUntilResourceRecordSetsChanged).toHaveBeenCalledTimes(1);
+  expect(mockWaitUntilResourceRecordSetsChanged.mock.calls[0][1]).toEqual({
     Id: 'change-id',
   });
 });

@@ -5,29 +5,17 @@ let mockDeleteParameters: jest.Mock ;
 let mockAddTagsToResource: jest.Mock;
 let mockGetParametersByPath: jest.Mock;
 let mockRemoveTagsFromResource: jest.Mock;
-jest.mock('aws-sdk', () => {
+
+jest.mock('@aws-sdk/client-ssm', () => {
   return {
-    SSM: jest.fn(() => {
-      return {
-        addTagsToResource: jest.fn((params) => {
-          return {
-            promise: () => mockAddTagsToResource(params),
-          };
-        }),
-        removeTagsFromResource: jest.fn((params) => {
-          return {
-            promise: () => mockRemoveTagsFromResource(params),
-          };
-        }),
-        getParametersByPath: jest.fn((params) => {
-          return {
-            promise: () => mockGetParametersByPath(params),
-          };
-        }),
-      };
-    }),
+    SSM: jest.fn().mockImplementation(() => ({
+      addTagsToResource: mockAddTagsToResource,
+      removeTagsFromResource: mockRemoveTagsFromResource,
+      getParametersByPath: mockGetParametersByPath,
+    })),
   };
 });
+
 beforeEach(() => {
   jest.spyOn(console, 'info').mockImplementation(() => {});
   jest.spyOn(console, 'error').mockImplementation(() => {});

diff --git a/packages/aws-cdk-lib/custom-resources/test/provider-framework/outbound.test.ts b/packages/aws-cdk-lib/custom-resources/test/provider-framework/outbound.test.ts
@@ -1,23 +1,21 @@
-import { GetFunctionResponse, InvokeCommandInput } from '@aws-sdk/client-lambda';
-import * as aws from 'aws-sdk';
+import { InvocationResponse, InvokeCommandInput } from '@aws-sdk/client-lambda';
 import { invokeFunction } from '../../lib/provider-framework/runtime/outbound';
 
-jest.mock('aws-sdk', () => {
-  return {
-    Lambda: class {
-      public invoke() {
-        return { promise: () => mockInvoke() };
-      }
+const mockInvoke = jest.fn();
+const mockWaitUntilFunctionActive = jest.fn();
+const mockLambda = {
+  invoke: mockInvoke,
+};
 
-      public waitFor() {
-        return { promise: () => mockWaitFor() };
-      }
-    },
+jest.mock('@aws-sdk/client-lambda', () => {
+  return {
+    Lambda: jest.fn().mockImplementation(() => {
+      return mockLambda;
+    }),
+    waitUntilFunctionActive: (...args: any[]) => mockWaitUntilFunctionActive(...args),
   };
 });
 
-let mockInvoke: () => Promise<aws.Lambda.InvocationResponse>;
-
 const req: InvokeCommandInput = {
   FunctionName: 'Whatever',
   Payload: JSON.stringify({
@@ -30,7 +28,7 @@ let invokeCount: number = 0;
 let expectedFunctionStates: string[] = [];
 let receivedFunctionStates: string[] = [];
 
-const mockWaitFor = async (): Promise<GetFunctionResponse> => {
+mockWaitUntilFunctionActive.mockImplementation(async () => {
   let state = expectedFunctionStates.pop();
   while (state !== 'Active') {
     receivedFunctionStates.push(state!);
@@ -48,7 +46,7 @@ const mockWaitFor = async (): Promise<GetFunctionResponse> => {
       State: 'Active',
     },
   };
-};
+});
 
 describe('invokeFunction tests', () => {
   afterEach(() => {
@@ -59,14 +57,14 @@ describe('invokeFunction tests', () => {
 
   // Success cases
   test('Inactive function that reactivates does not throw error', async () => {
-    mockInvoke = async () => {
+    mockInvoke.mockImplementation(async () => {
       if (invokeCount == 0) {
         invokeCount++;
         throw new Error('Better luck next time');
       }
       invokeCount++;
-      return { Payload: req.Payload };
-    };
+      return { Payload: req.Payload } as InvocationResponse;
+    });
 
     expectedFunctionStates.push('Active');
     expectedFunctionStates.push('Pending');
@@ -77,14 +75,14 @@ describe('invokeFunction tests', () => {
   });
 
   test('Active function does not run waitFor or retry invoke', async () => {
-    mockInvoke = async () => {
+    mockInvoke.mockImplementation(async () => {
       if (invokeCount == 1) {
         invokeCount++;
         throw new Error('This should not happen in this test');
       }
       invokeCount++;
-      return { Payload: req.Payload };
-    };
+      return { Payload: req.Payload } as InvocationResponse;
+    });
 
     expectedFunctionStates.push('Active');
 
@@ -95,14 +93,14 @@ describe('invokeFunction tests', () => {
 
   // Failure cases
   test('Inactive function that goes back to inactive throws error', async () => {
-    mockInvoke = async () => {
+    mockInvoke.mockImplementation(async () => {
       if (invokeCount == 0) {
         invokeCount++;
         throw new Error('Better luck next time');
       }
       invokeCount++;
-      return { Payload: req.Payload };
-    };
+      return { Payload: req.Payload } as InvocationResponse;
+    });
 
     expectedFunctionStates.push('Inactive');
     expectedFunctionStates.push('Pending');
@@ -114,14 +112,14 @@ describe('invokeFunction tests', () => {
   });
 
   test('Inactive function that goes to failed throws error', async () => {
-    mockInvoke = async () => {
+    mockInvoke.mockImplementation(async () => {
       if (invokeCount == 0) {
         invokeCount++;
         throw new Error('Better luck next time');
       }
       invokeCount++;
-      return { Payload: req.Payload };
-    };
+      return { Payload: req.Payload } as InvocationResponse;
+    });
 
     expectedFunctionStates.push('Failed');
     expectedFunctionStates.push('Pending');
@@ -133,14 +131,14 @@ describe('invokeFunction tests', () => {
   });
 
   test('Inactive function that returns other value throws error', async () => {
-    mockInvoke = async () => {
+    mockInvoke.mockImplementation(async () => {
       if (invokeCount == 0) {
         invokeCount++;
         throw new Error('Better luck next time');
       }
       invokeCount++;
-      return { Payload: req.Payload };
-    };
+      return { Payload: req.Payload } as InvocationResponse;
+    });
 
     expectedFunctionStates.push('NewFunctionWhoDis');
     expectedFunctionStates.push('Pending');
@@ -152,14 +150,14 @@ describe('invokeFunction tests', () => {
   });
 
   test('Wait for stops on terminal responses', async () => {
-    mockInvoke = async () => {
+    mockInvoke.mockImplementation(async () => {
       if (invokeCount == 0) {
         invokeCount++;
         throw new Error('Better luck next time');
       }
       invokeCount++;
-      return { Payload: req.Payload };
-    };
+      return { Payload: req.Payload } as InvocationResponse;
+    });
 
     expectedFunctionStates.push('SomethingElse');
     expectedFunctionStates.push('Pending');