Upgrade scale up/down lambdas to aws sdk v3 (#7061)

Upgrade to aws sdk v3

Main change is getting rid of the `promise` calls since I think they
just directly return promises instead of requests

This changes a lot of mocks in the testing so I'm not sure how good
running just `yarn test` is

Testing:
Mangled scaleDown to only run `listInstances` and `listSSMParameters`
In `terraform-aws-github-runner/modules/runners/lambdas/runners`:
```
yarn build; cd dist;node -e 'require("./index").scaleDown({}, {}, {});' > t.log    
```
I also tried to terminate a runner and it worked


Deployed to pytorch-canary and ran some jobs, seems ok

Author: clee2000

terraform-aws-github-runner/modules/runners/lambdas/runners/jest.config.js

@@ -12,7 +12,4 @@ module.exports = {
       statements: 94
     }
   },
-  moduleNameMapper: {
-    axios: 'axios/dist/node/axios.cjs', // Allow axios to work in tests
-  },
 };

terraform-aws-github-runner/modules/runners/lambdas/runners/package.json

@@ -18,20 +18,26 @@
   "devDependencies": {
     "@types/aws-lambda": "^8.10.72",
     "@types/express": "^4.17.11",
-    "@types/jest": "^26.0.20",
+    "@types/jest": "29",
     "@typescript-eslint/eslint-plugin": "^5.30.5",
     "@typescript-eslint/parser": "^5.30.5",
     "@vercel/ncc": "^0.38.1",
     "eslint": "^8.19.0",
-    "jest": "^26.6.3",
+    "jest": "29",
     "jest-mock-extended": "^1.0.13",
     "moment-timezone": "^0.5.35",
     "nock": "^13.0.11",
     "prettier": "^2.4.1",
-    "ts-jest": "^26.5.3",
+    "ts-jest": "29",
     "ts-node-dev": "^1.1.6"
   },
   "dependencies": {
+    "@aws-sdk/client-cloudwatch": "^3.876.0",
+    "@aws-sdk/client-ec2": "^3.876.0",
+    "@aws-sdk/client-kms": "^3.876.0",
+    "@aws-sdk/client-secrets-manager": "^3.876.0",
+    "@aws-sdk/client-sqs": "^3.876.0",
+    "@aws-sdk/client-ssm": "^3.876.0",
     "@octokit/auth-app": "^3.0.0",
     "@octokit/request-error": "^2.1.0",
     "@octokit/rest": "^18.3.5",
@@ -41,7 +47,6 @@
     "@types/node": "^14.14.34",
     "@types/uuid": "^9.0.1",
     "async-mutex": "^0.4.0",
-    "aws-sdk": "^2.863.0",
     "axios": "^1.8.2",
     "cron-parser": "^3.3.0",
     "generic-pool": "^3.9.0",

terraform-aws-github-runner/modules/runners/lambdas/runners/src/lambda.test.ts

@@ -3,19 +3,17 @@ import { scaleDown as scaleDownL, scaleUp as scaleUpL, scaleUpChron as scaleUpCh
 import nock from 'nock';
 import { Config } from './scale-runners/config';
 import { Context, SQSEvent, ScheduledEvent } from 'aws-lambda';
-import { mocked } from 'ts-jest/utils';
+import { mocked } from 'jest-mock';
 import { scaleDown } from './scale-runners/scale-down';
 import { scaleUp, RetryableScalingError } from './scale-runners/scale-up';
 import { scaleUpChron } from './scale-runners/scale-up-chron';
 import { sqsSendMessages, sqsDeleteMessageBatch } from './scale-runners/sqs';
 import * as MetricsModule from './scale-runners/metrics';
 
 const mockCloudWatch = {
-  putMetricData: jest.fn().mockImplementation(() => {
-    return { promise: jest.fn().mockResolvedValue(true) };
-  }),
+  putMetricData: jest.fn().mockResolvedValue(true),
 };
-jest.mock('aws-sdk', () => ({
+jest.mock('@aws-sdk/client-cloudwatch', () => ({
   CloudWatch: jest.fn().mockImplementation(() => mockCloudWatch),
 }));
 

terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/cache.test.ts

@@ -7,7 +7,7 @@ import {
   getExperimentValue,
   getJoinedStressTestExperiment,
 } from './cache';
-import { mocked } from 'ts-jest/utils';
+import { mocked } from 'jest-mock';
 import { v4 as uuidv4 } from 'uuid';
 import nock from 'nock';
 import { RedisClientType, createClient } from 'redis';
@@ -47,7 +47,7 @@ beforeEach(() => {
   jest.restoreAllMocks();
   nock.disableNetConnect();
 
-  mocked(createClient).mockImplementation(produceMockedRedis);
+  (mocked(createClient) as any).mockImplementation(produceMockedRedis);
 
   jest.spyOn(Config, 'Instance', 'get').mockImplementation(() => config as unknown as Config);
 });

terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/gh-auth.test.ts

@@ -28,13 +28,12 @@ jest.mock('./cache', () => ({
     }),
 }));
 
-const mockSMgetSecretValuePromise = jest.fn();
 const mockSMgetSecretValue = jest.fn();
-jest.mock('aws-sdk', () => ({
+
+jest.mock('@aws-sdk/client-secrets-manager', () => ({
   SecretsManager: jest.fn().mockImplementation(() => ({
     getSecretValue: mockSMgetSecretValue,
   })),
-  CloudWatch: jest.requireActual('aws-sdk').CloudWatch,
 }));
 
 const metrics = new ScaleUpMetrics();
@@ -95,8 +94,7 @@ describe('Test createGithubAuth', () => {
   describe('tests where aws-sdk fails', () => {
     const message = 'Error message on exception';
     beforeEach(() => {
-      mockSMgetSecretValuePromise.mockClear().mockRejectedValue(Error(message));
-      mockSMgetSecretValue.mockClear().mockImplementation(() => ({ promise: mockSMgetSecretValuePromise }));
+      mockSMgetSecretValue.mockClear().mockRejectedValue(Error(message));
 
       jest.spyOn(Config, 'Instance', 'get').mockImplementation(
         () =>
@@ -115,11 +113,10 @@ describe('Test createGithubAuth', () => {
 
   describe('tests where aws-sdk works as expected', () => {
     beforeEach(() => {
-      mockSMgetSecretValuePromise
+      mockSMgetSecretValue
         .mockClear()
         .mockResolvedValueOnce({ SecretString: undefined })
         .mockResolvedValueOnce({ SecretString: secretString });
-      mockSMgetSecretValue.mockClear().mockImplementation(() => ({ promise: mockSMgetSecretValuePromise }));
     });
 
     describe('github keys are not from environment, nor secretsManagerSecretsId is provided', () => {
@@ -286,12 +283,12 @@ describe('Test createGithubAuth', () => {
         const result2 = await createGithubAuth(installationId, authType, '', metrics);
 
         expect(mockSMgetSecretValue).toBeCalledTimes(2);
-        expect(mockSMgetSecretValue).toHaveBeenCalledWith({ SecretId: Config.Instance.secretsManagerSecretsId });
-        expect(mockSMgetSecretValuePromise).toBeCalledTimes(2);
+        expect(mockSMgetSecretValue).toHaveBeenCalledWith({
+          SecretId: Config.Instance.secretsManagerSecretsId,
+        });
 
         expect(mockedDecrypt).toBeCalledWith('github_app_client_secret', config.kmsKeyId, config.environment, metrics);
         expect(mockedDecrypt).toBeCalledWith('github_app_key_base64', config.kmsKeyId, config.environment, metrics);
-        expect(mockSMgetSecretValuePromise).toBeCalledTimes(2);
         expect(mockedCreatAppAuth).toBeCalledWith(authOptions);
         expect(mockedAuth).toBeCalledWith({ type: authType });
         expect(result1).toBe(token);

terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/gh-auth.ts

@@ -2,7 +2,7 @@ import { Config } from './config';
 import { Metrics } from './metrics';
 import { Octokit } from '@octokit/rest';
 import { OctokitOptions } from '@octokit/core/dist-types/types';
-import { SecretsManager } from 'aws-sdk';
+import { SecretsManager } from '@aws-sdk/client-secrets-manager';
 import {
   AppAuthentication,
   InstallationAccessTokenAuthentication,
@@ -39,7 +39,7 @@ async function getCredentialsFromSecretsManager(
           metrics.smGetSecretValueAWSCallSuccess,
           metrics.smGetSecretValueAWSCallFailure,
           () => {
-            return secretsManager.getSecretValue({ SecretId: secretsManagerSecretsId }).promise();
+            return secretsManager.getSecretValue({ SecretId: secretsManagerSecretsId });
           },
         );
       });

terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/gh-issues.test.ts

@@ -2,7 +2,7 @@ import { getRepoIssuesWithLabel, resetIssuesCache } from './gh-issues';
 
 import { Octokit } from '@octokit/rest';
 import { createGitHubClientForRunnerRepo } from './gh-runners';
-import { mocked } from 'ts-jest/utils';
+import { mocked } from 'jest-mock';
 import nock from 'nock';
 import { ScaleUpMetrics } from './metrics';
 import { redisCached } from './cache';

terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/gh-runners.test.ts

@@ -20,26 +20,9 @@ import { ScaleUpMetrics } from './metrics';
 
 import { Config } from './config';
 import { Octokit } from '@octokit/rest';
-import { mocked } from 'ts-jest/utils';
+import { mocked } from 'jest-mock';
 import { locallyCached, redisCached } from './cache';
 
-const mockEC2 = {
-  describeInstances: jest.fn(),
-  runInstances: jest.fn(),
-  terminateInstances: jest.fn().mockReturnValue({ promise: jest.fn() }),
-};
-const mockSSMdescribeParametersRet = jest.fn();
-const mockSSM = {
-  deleteParameter: jest.fn().mockReturnValue({ promise: jest.fn() }),
-  describeParameters: jest.fn().mockReturnValue({ promise: mockSSMdescribeParametersRet }),
-  putParameter: jest.fn().mockReturnValue({ promise: jest.fn() }),
-};
-jest.mock('aws-sdk', () => ({
-  EC2: jest.fn().mockImplementation(() => mockEC2),
-  SSM: jest.fn().mockImplementation(() => mockSSM),
-  CloudWatch: jest.requireActual('aws-sdk').CloudWatch,
-}));
-
 jest.mock('./gh-auth');
 jest.mock('./cache', () => ({
   /* eslint-disable-next-line @typescript-eslint/no-explicit-any */

terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/kms/index.test.ts

@@ -1,4 +1,3 @@
-import AWS from 'aws-sdk';
 import { Config } from '../config';
 import { decrypt } from './index';
 import { ScaleUpMetrics } from '../metrics';
@@ -11,22 +10,12 @@ const mockKmsPromise = {
     toString: jest.fn().mockReturnValue(decryptedStr),
   },
 };
-const mockKmsDecrypt = {
-  promise: jest.fn().mockImplementation(async () => mockKmsPromise),
-};
 const mockKms = {
-  decrypt: jest.fn().mockImplementation(() => mockKmsDecrypt),
+  decrypt: jest.fn().mockResolvedValue(mockKmsPromise),
 };
 
-jest.mock('aws-sdk', () => ({
-  __esModule: true,
-  default: {
-    config: {
-      update: jest.fn(),
-    },
-  },
+jest.mock('@aws-sdk/client-kms', () => ({
   KMS: jest.fn().mockImplementation(() => mockKms),
-  CloudWatch: jest.requireActual('aws-sdk').CloudWatch,
 }));
 
 beforeEach(() => {
@@ -37,7 +26,7 @@ beforeEach(() => {
 });
 
 describe('decrypt', () => {
-  describe('check AWS && KMS calls', () => {
+  describe('check KMS calls', () => {
     it('simple calls decrypt', async () => {
       const encrypted = Buffer.from('some random buffer');
       const key = 'decrypt key';
@@ -53,17 +42,13 @@ describe('decrypt', () => {
       expect(await decrypt(encrypted.toString('base64'), key, environmentName, new ScaleUpMetrics())).toBe(
         decryptedStr,
       );
-      expect(AWS.config.update).toBeCalledWith({
-        region: awsRegion,
-      });
       expect(mockKms.decrypt).toBeCalledWith({
         CiphertextBlob: encrypted,
         KeyId: key,
         EncryptionContext: {
           ['Environment']: environmentName,
         },
       });
-      expect(mockKmsDecrypt.promise).toBeCalled();
       expect(mockKmsPromise.Plaintext.toString).toBeCalled();
     });
   });

terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/kms/index.ts

@@ -1,7 +1,6 @@
-import AWS from 'aws-sdk';
+import { KMS } from '@aws-sdk/client-kms';
 import { Config } from '../config';
 import { expBackOff } from '../utils';
-import { KMS } from 'aws-sdk';
 import { Metrics } from '../metrics';
 
 let kms: KMS | undefined = undefined;
@@ -14,27 +13,23 @@ export async function decrypt(
 ): Promise<string | undefined> {
   /* istanbul ignore next */
   if (!kms) {
-    AWS.config.update({
+    kms = new KMS({
       region: Config.Instance.awsRegion,
     });
-
-    kms = new KMS();
   }
 
   // this is so the linter understands that KMS is not undefined at this point :(
   const kmsD = kms;
 
   const decripted = await expBackOff(() => {
     return metrics.trackRequest(metrics.kmsDecryptAWSCallSuccess, metrics.kmsDecryptAWSCallFailure, () => {
-      return kmsD
-        .decrypt({
-          CiphertextBlob: Buffer.from(encrypted, 'base64'),
-          KeyId: key,
-          EncryptionContext: {
-            ['Environment']: environmentName,
-          },
-        })
-        .promise();
+      return kmsD.decrypt({
+        CiphertextBlob: Buffer.from(encrypted, 'base64'),
+        KeyId: key,
+        EncryptionContext: {
+          ['Environment']: environmentName,
+        },
+      });
     });
   });