refactor, prepare for creating muliple runners

Author: npalm

modules/runners/lambdas/runners/jest.config.js

@@ -5,10 +5,10 @@ module.exports = {
   collectCoverageFrom: ['src/**/*.{ts,js,jsx}', '!src/**/*local*.ts', 'src/**/*.d.ts'],
   coverageThreshold: {
     global: {
-      branches: 80,
-      functions: 80,
-      lines: 80,
-      statements: 80,
+      branches: 92,
+      functions: 92,
+      lines: 92,
+      statements: 92,
     },
   },
 };

modules/runners/lambdas/runners/src/aws/runners.test.ts

@@ -126,6 +126,25 @@ describe('list instances', () => {
     const resp = await listEC2Runners();
     expect(resp.length).toBe(0);
   });
+
+  it('Instances with no tags.', async () => {
+    const noInstances: AWS.EC2.DescribeInstancesResult = {
+      Reservations: [
+        {
+          Instances: [
+            {
+              LaunchTime: new Date('2020-10-11T14:48:00.000+09:00'),
+              InstanceId: 'i-5678',
+              Tags: undefined,
+            },
+          ],
+        },
+      ],
+    };
+    mockDescribeInstances.promise.mockReturnValue(noInstances);
+    const resp = await listEC2Runners();
+    expect(resp.length).toBe(1);
+  });
 });
 
 describe('terminate runner', () => {
@@ -159,6 +178,7 @@ describe('create runner', () => {
     type: 'Org',
     capacityType: 'spot',
     allocationStrategy: 'capacity-optimized',
+    totalTargetCapacity: 1,
   };
 
   beforeEach(() => {
@@ -172,23 +192,47 @@ describe('create runner', () => {
     mockSSM.putParameter.mockImplementation(() => mockPutParameter);
   });
 
-  it('calls run instances with the correct config for repo', async () => {
+  it('calls create fleet of 1 instance with the correct config for repo', async () => {
     await createRunner(createRunnerConfig({ ...defaultRunnerConfig, type: 'Repo' }));
     expect(mockEC2.createFleet).toBeCalledWith(
       expectedCreateFleetRequest({ ...defaultExpectedFleetRequestValues, type: 'Repo' }),
     );
+    expect(mockSSM.putParameter).toBeCalledTimes(1);
+  });
+
+  it('calls create fleet of 2 instances with the correct config for org ', async () => {
+    const instances = [{ InstanceIds: ['i-1234', 'i-5678'] }];
+    mockCreateFleet.promise.mockReturnValue({
+      Instances: instances,
+    });
+
+    await createRunner({ ...createRunnerConfig(defaultRunnerConfig), numberOfRunners: 2 });
+
+    expect(mockEC2.createFleet).toBeCalledWith(
+      expectedCreateFleetRequest({ ...defaultExpectedFleetRequestValues, totalTargetCapacity: 2 }),
+    );
+    expect(mockSSM.putParameter).toBeCalledTimes(2);
+    for (const instance of instances[0].InstanceIds) {
+      expect(mockSSM.putParameter).toBeCalledWith({
+        Name: `unit-test-environment-${instance}`,
+        Type: 'SecureString',
+        Value: 'bla',
+      });
+    }
   });
 
-  it('calls run instances with the correct config for org', async () => {
+  it('calls create fleet of 1 instance with the correct config for org', async () => {
     await createRunner(createRunnerConfig(defaultRunnerConfig));
     expect(mockEC2.createFleet).toBeCalledWith(expectedCreateFleetRequest(defaultExpectedFleetRequestValues));
+    expect(mockSSM.putParameter).toBeCalledTimes(1);
   });
 
-  it('calls run instances with the on-demand capacity', async () => {
+  it('calls create fleet of 1 instance with the on-demand capacity', async () => {
     await createRunner(createRunnerConfig({ ...defaultRunnerConfig, capacityType: 'on-demand' }));
     expect(mockEC2.createFleet).toBeCalledWith(
       expectedCreateFleetRequest({ ...defaultExpectedFleetRequestValues, capacityType: 'on-demand' }),
     );
+    expect(mockSSM.putParameter).toBeCalledTimes(1);
   });
 
   it('calls run instances with the on-demand capacity', async () => {
@@ -226,6 +270,7 @@ describe('create runner with errors', () => {
     type: 'Repo',
     capacityType: 'spot',
     allocationStrategy: 'capacity-optimized',
+    totalTargetCapacity: 1,
   };
   beforeEach(() => {
     jest.clearAllMocks();
@@ -332,6 +377,7 @@ interface ExpectedFleetRequestValues {
   capacityType: EC2.DefaultTargetCapacityType;
   allocationStrategy: EC2.AllocationStrategy;
   maxSpotPrice?: string;
+  totalTargetCapacity: number;
 }
 
 function expectedCreateFleetRequest(expectedValues: ExpectedFleetRequestValues): AWS.EC2.CreateFleetRequest {
@@ -378,7 +424,7 @@ function expectedCreateFleetRequest(expectedValues: ExpectedFleetRequestValues):
     ],
     TargetCapacitySpecification: {
       DefaultTargetCapacityType: expectedValues.capacityType,
-      TotalTargetCapacity: 1,
+      TotalTargetCapacity: expectedValues.totalTargetCapacity,
     },
     Type: 'instant',
   };

modules/runners/lambdas/runners/src/aws/runners.ts

@@ -39,6 +39,7 @@ export interface RunnerInputParameters {
     maxSpotPrice?: string;
     instanceAllocationStrategy: EC2.SpotAllocationStrategy;
   };
+  numberOfRunners?: number;
 }
 
 export async function listEC2Runners(filters: ListRunnerFilters | undefined = undefined): Promise<RunnerList[]> {
@@ -107,6 +108,7 @@ export async function createRunner(runnerParameters: RunnerInputParameters): Pro
   logger.debug('Runner configuration: ' + JSON.stringify(runnerParameters), LogFields.print());
 
   const ec2 = new EC2();
+  const numberOfRunners = runnerParameters.numberOfRunners ? runnerParameters.numberOfRunners : 1;
 
   let fleet: AWS.EC2.CreateFleetResult;
   try {
@@ -130,7 +132,7 @@ export async function createRunner(runnerParameters: RunnerInputParameters): Pro
           AllocationStrategy: 'capacity-optimized',
         },
         TargetCapacitySpecification: {
-          TotalTargetCapacity: 1,
+          TotalTargetCapacity: numberOfRunners,
           DefaultTargetCapacityType: runnerParameters.ec2instanceCriteria.targetCapacityType,
         },
         TagSpecifications: [

modules/runners/lambdas/runners/src/aws/ssm.test.ts

@@ -37,4 +37,22 @@ describe('Test getParameterValue', () => {
     // Assert
     expect(result).toBe(parameterValue);
   });
+
+  test('Gets invalid parameters and returns string', async () => {
+    // Arrange
+    const parameterName = 'invalid';
+    const output: GetParameterCommandOutput = {
+      $metadata: {
+        httpStatusCode: 200,
+      },
+    };
+
+    SSM.prototype.getParameter = jest.fn().mockResolvedValue(output);
+
+    // Act
+    const result = await getParameterValue(parameterName);
+
+    // Assert
+    expect(result).toBe(undefined);
+  });
 });

modules/runners/lambdas/runners/src/scale-runners/scale-up.test.ts

@@ -379,6 +379,14 @@ describe('scaleUp with public GH', () => {
     expect(listEC2Runners).not.toBeCalled();
   });
 
+  it('does not list runners when no workflows are queued (check_run)', async () => {
+    mockOctokit.checks.get.mockImplementation(() => ({
+      data: { status: 'completed' },
+    }));
+    await scaleUpModule.scaleUp('aws:sqs', { ...TEST_DATA, eventType: 'check_run' });
+    expect(listEC2Runners).not.toBeCalled();
+  });
+
   describe('on org level', () => {
     beforeEach(() => {
       process.env.ENABLE_ORGANIZATION_RUNNERS = 'true';

modules/runners/lambdas/runners/src/scale-runners/scale-up.ts

@@ -1,4 +1,4 @@
-import { listEC2Runners, createRunner } from './../aws/runners';
+import { listEC2Runners, createRunner, RunnerInputParameters } from './../aws/runners';
 import { createOctoClient, createGithubAppAuth, createGithubInstallationAuth } from '../gh-auth/gh-auth';
 import yn from 'yn';
 import { Octokit } from '@octokit/rest';
@@ -95,6 +95,44 @@ async function isJobQueued(githubInstallationClient: Octokit, payload: ActionReq
   return isQueued;
 }
 
+async function createRunners(
+  enableOrgLevel: boolean,
+  githubInstallationClient: Octokit,
+  payload: ActionRequestMessage,
+  runnerExtraLabels: string | undefined,
+  runnerGroup: string | undefined,
+  ghesBaseUrl: string,
+  ephemeral: boolean,
+  runnerType: 'Org' | 'Repo',
+  environment: string,
+  runnerOwner: string,
+  subnets: string[],
+  launchTemplateName: string,
+  ec2instanceCriteria: RunnerInputParameters['ec2instanceCriteria'],
+): Promise<void> {
+  const token = await getGithubRunnerRegistrationToken(enableOrgLevel, githubInstallationClient, payload);
+
+  const runnerServiceConfig = generateRunnerServiceConfig(
+    runnerExtraLabels,
+    runnerGroup,
+    ghesBaseUrl,
+    ephemeral,
+    token,
+    runnerType,
+    payload,
+  );
+
+  await createRunner({
+    environment,
+    runnerServiceConfig,
+    runnerOwner,
+    runnerType,
+    subnets,
+    launchTemplateName,
+    ec2instanceCriteria,
+  });
+}
+
 export async function scaleUp(eventSource: string, payload: ActionRequestMessage): Promise<void> {
   logger.info(
     `Received ${payload.eventType} from ${payload.repositoryOwner}/${payload.repositoryName}`,
@@ -158,32 +196,26 @@ export async function scaleUp(eventSource: string, payload: ActionRequestMessage
     if (currentRunners.length < maximumRunners) {
       logger.info(`Attempting to launch a new runner`, LogFields.print());
 
-      const token = await getGithubRunnerRegistrationToken(enableOrgLevel, githubInstallationClient, payload);
-
-      const runnerServiceConfig = generateRunnerServiceConfig(
+      await createRunners(
+        enableOrgLevel,
+        githubInstallationClient,
+        payload,
         runnerExtraLabels,
         runnerGroup,
         ghesBaseUrl,
         ephemeral,
-        token,
         runnerType,
-        payload,
-      );
-
-      await createRunner({
         environment,
-        runnerServiceConfig,
         runnerOwner,
-        runnerType,
         subnets,
         launchTemplateName,
-        ec2instanceCriteria: {
+        {
           instanceTypes,
           targetCapacityType: instanceTargetTargetCapacityType,
           maxSpotPrice: instanceMaxSpotPrice,
           instanceAllocationStrategy: instanceAllocationStrategy,
         },
-      });
+      );
     } else {
       logger.info('No runner will be created, maximum number of runners reached.', LogFields.print());
       if (ephemeral) {