diff --git a/lambdas/functions/control-plane/src/pool/pool.test.ts b/lambdas/functions/control-plane/src/pool/pool.test.ts index a3e199d7c1..8372ab6403 100644 --- a/lambdas/functions/control-plane/src/pool/pool.test.ts +++ b/lambdas/functions/control-plane/src/pool/pool.test.ts @@ -1,51 +1,31 @@ -import { Octokit } from '@octokit/rest'; -import moment from 'moment-timezone'; -import * as nock from 'nock'; +import type { Octokit } from '@octokit/rest'; +import { defaultComputeProvider } from '@aws-github-runner/compute-providers/provider-types'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; -import { createRunners } from '@aws-github-runner/compute-providers/aws/ec2/control-plane/runner-config'; -import { listEC2Runners } from '@aws-github-runner/compute-providers/aws/ec2/runners'; import * as ghAuth from '../github/auth'; -import { getGitHubEnterpriseApiUrl } from '../scale-runners/github-runner'; +import { controlPlaneProviderRegistry } from '../control-plane-providers'; +import * as githubRunner from '../scale-runners/github-runner'; import { adjust } from './pool'; -import { describe, it, expect, beforeEach, vi, MockedClass } from 'vitest'; +import type { PoolComputeProvider } from './pool-provider'; -const mockOctokit = { +const githubClient = { paginate: vi.fn(), - checks: { get: vi.fn() }, actions: { - createRegistrationTokenForOrg: vi.fn(), + listSelfHostedRunnersForOrg: vi.fn(), }, apps: { getOrgInstallation: vi.fn(), }, }; -vi.mock('@octokit/rest', () => ({ - Octokit: vi.fn().mockImplementation(function () { - return mockOctokit; - }), -})); - -vi.mock('@aws-github-runner/compute-providers/aws/ec2/runners', async () => ({ - listEC2Runners: vi.fn(), - // Include any other functions from the module that might be used - bootTimeExceeded: vi.fn(), -})); -vi.mock('./../github/auth', async () => ({ +vi.mock('../github/auth', () => ({ createGithubAppAuth: vi.fn(), createGithubInstallationAuth: vi.fn(), createOctokitClient: vi.fn(), getStoredInstallationId: vi.fn().mockResolvedValue(undefined), })); -vi.mock('@aws-github-runner/compute-providers/aws/ec2/control-plane/runner-config', async (importOriginal) => ({ - ...(await importOriginal< - typeof import('@aws-github-runner/compute-providers/aws/ec2/control-plane/runner-config') - >()), - createRunners: vi.fn(), -})); - -vi.mock('../scale-runners/github-runner', async () => ({ +vi.mock('../scale-runners/github-runner', () => ({ createStartRunnerConfig: vi.fn(), getGitHubEnterpriseApiUrl: vi.fn().mockReturnValue({ ghesApiUrl: '', @@ -54,50 +34,28 @@ vi.mock('../scale-runners/github-runner', async () => ({ validateSsmParameterStoreTags: vi.fn().mockReturnValue([]), })); -const mocktokit = Octokit as MockedClass; const mockedAppAuth = vi.mocked(ghAuth.createGithubAppAuth); const mockedInstallationAuth = vi.mocked(ghAuth.createGithubInstallationAuth); -const mockCreateClient = vi.mocked(ghAuth.createOctokitClient); -const mockListRunners = vi.mocked(listEC2Runners); +const mockedCreateClient = vi.mocked(ghAuth.createOctokitClient); +const mockedResolveCapability = vi.spyOn(controlPlaneProviderRegistry, 'capability'); +const mockedGetGitHubEnterpriseApiUrl = vi.mocked(githubRunner.getGitHubEnterpriseApiUrl); + +const poolProvider = { + listRunners: vi.fn(), + countAvailableRunners: vi.fn(), + createRunners: vi.fn(), +} satisfies Omit; const cleanEnv = process.env; const ORG = 'my-org'; -const MINIMUM_TIME_RUNNING = 15; -const ec2InstancesRegistered = [ - { - id: 'i-1-idle', - launchTime: new Date(), - type: 'Org', - owner: ORG, - }, - { - id: 'i-2-busy', - launchTime: new Date(), - type: 'Org', - owner: ORG, - }, - { - id: 'i-3-offline', - launchTime: new Date(), - type: 'Org', - owner: ORG, - }, - { - id: 'i-4-idle-older-than-minimum-time-running', - launchTime: moment(new Date()) - .subtract(MINIMUM_TIME_RUNNING + 3, 'minutes') - .toDate(), - type: 'Org', - owner: ORG, - }, -]; +const providerRunners = [{ id: 'runner-1' }, { id: 'runner-2' }, { id: 'runner-3' }, { id: 'runner-4' }]; const githubRunnersRegistered = [ { id: 1, - name: 'i-1-idle', + name: 'runner-1', os: 'linux', status: 'online', busy: false, @@ -105,7 +63,7 @@ const githubRunnersRegistered = [ }, { id: 2, - name: 'i-2-busy', + name: 'runner-2', os: 'linux', status: 'online', busy: true, @@ -113,15 +71,15 @@ const githubRunnersRegistered = [ }, { id: 3, - name: 'i-3-offline', + name: 'runner-3', os: 'linux', status: 'offline', busy: false, labels: [], }, { - id: 3, - name: 'i-4-idle-older-than-minimum-time-running', + id: 4, + name: 'runner-4', os: 'linux', status: 'online', busy: false, @@ -130,49 +88,21 @@ const githubRunnersRegistered = [ ]; beforeEach(() => { - nock.disableNetConnect(); - vi.resetModules(); vi.clearAllMocks(); process.env = { ...cleanEnv }; - process.env.GITHUB_APP_KEY_BASE64 = 'TEST_CERTIFICATE_DATA'; - process.env.GITHUB_APP_ID = '1337'; - process.env.GITHUB_APP_CLIENT_ID = 'TEST_CLIENT_ID'; - process.env.GITHUB_APP_CLIENT_SECRET = 'TEST_CLIENT_SECRET'; process.env.RUNNERS_MAXIMUM_COUNT = '-1'; process.env.ENVIRONMENT = 'unit-test-environment'; - process.env.ENABLE_ORGANIZATION_RUNNERS = 'true'; - process.env.LAUNCH_TEMPLATE_NAME = 'lt-1'; - process.env.SUBNET_IDS = 'subnet-123'; process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/tokens'; - process.env.INSTANCE_TYPES = 'm5.large'; - process.env.INSTANCE_TARGET_CAPACITY_TYPE = 'spot'; process.env.RUNNER_OWNER = ORG; - process.env.RUNNER_BOOT_TIME_IN_MINUTES = MINIMUM_TIME_RUNNING.toString(); - process.env.SCALE_ERRORS = - '["UnfulfillableCapacity","MaxSpotInstanceCountExceeded","TargetCapacityLimitExceededException"]'; - - const mockTokenReturnValue = { - data: { - token: '1234abcd', - }, - }; - mockOctokit.actions.createRegistrationTokenForOrg.mockImplementation(() => mockTokenReturnValue); - - mockOctokit.paginate.mockImplementation(() => githubRunnersRegistered); - - mockListRunners.mockImplementation(async () => ec2InstancesRegistered); - vi.mocked(createRunners).mockResolvedValue({ - instances: [], - retryableErrorCount: 0, - nonRetryableErrorCount: 0, - }); - const mockInstallationIdReturnValueOrgs = { - data: { - id: 1, - }, - }; - mockOctokit.apps.getOrgInstallation.mockImplementation(() => mockInstallationIdReturnValueOrgs); + githubClient.paginate.mockResolvedValue(githubRunnersRegistered); + githubClient.apps.getOrgInstallation.mockResolvedValue({ data: { id: 1 } }); + mockedGetGitHubEnterpriseApiUrl.mockReturnValue({ ghesApiUrl: '', ghesBaseUrl: '' }); + + mockedResolveCapability.mockReturnValue(() => poolProvider); + poolProvider.listRunners.mockResolvedValue(providerRunners); + poolProvider.countAvailableRunners.mockReturnValue(2); + poolProvider.createRunners.mockResolvedValue([]); mockedAppAuth.mockResolvedValue({ type: 'app', @@ -192,169 +122,92 @@ beforeEach(() => { installationId: 0, }); - mockCreateClient.mockResolvedValue(new mocktokit()); + mockedCreateClient.mockResolvedValue(githubClient as unknown as Octokit); }); -describe('Test simple pool.', () => { +describe('pool adjustment', () => { describe('With GitHub Cloud', () => { beforeEach(() => { - (getGitHubEnterpriseApiUrl as ReturnType).mockReturnValue({ + mockedGetGitHubEnterpriseApiUrl.mockReturnValue({ ghesApiUrl: '', ghesBaseUrl: '', }); }); - it('Top up pool with pool size 2 registered.', async () => { - await adjust({ poolSize: 3, type: 'ec2' }); - expect(createRunners).toHaveBeenCalledTimes(1); - expect(createRunners).toHaveBeenCalledWith( - expect.anything(), - expect.anything(), - 1, - expect.anything(), - expect.anything(), - 'pool-lambda', - ); + + it('tops up the pool to the requested size', async () => { + await adjust({ poolSize: 3 }); + + expect(poolProvider.createRunners).toHaveBeenCalledTimes(1); + expect(poolProvider.createRunners).toHaveBeenCalledWith(expect.objectContaining({ numberOfRunners: 1 })); }); - it('Defaults legacy pool events without a provider type to EC2.', async () => { + it('uses the default provider for events without a provider type', async () => { await adjust({ poolSize: 10 }); - expect(mockListRunners).toHaveBeenCalledWith({ + + expect(mockedResolveCapability).toHaveBeenCalledWith(defaultComputeProvider, 'pool'); + expect(poolProvider.listRunners).toHaveBeenCalledWith({ environment: 'unit-test-environment', runnerOwner: ORG, runnerType: 'Org', - statuses: ['running'], }); - expect(createRunners).toHaveBeenCalledWith( - expect.anything(), - expect.anything(), - 8, - expect.anything(), - expect.anything(), - 'pool-lambda', - ); - }); - - it('Selects the EC2 pool provider case-insensitively.', async () => { - await adjust({ poolSize: 10, type: ' EC2 ' }); - expect(createRunners).toHaveBeenCalledWith( - expect.anything(), - expect.anything(), - 8, - expect.anything(), - expect.anything(), - 'pool-lambda', - ); + expect(poolProvider.createRunners).toHaveBeenCalledWith(expect.objectContaining({ numberOfRunners: 8 })); }); - it('Rejects unsupported pool provider types.', async () => { + it('rejects unsupported provider types', async () => { await expect(adjust({ poolSize: 10, type: 'unsupported-provider' })).rejects.toThrow( "Unsupported compute provider type 'unsupported-provider'", ); - expect(mockListRunners).not.toHaveBeenCalled(); - }); - - it('Should not top up if pool size is reached.', async () => { - await adjust({ poolSize: 1, type: 'ec2' }); - expect(createRunners).not.toHaveBeenCalled(); - }); - - it('Should top up if pool size is not reached including a booting instance.', async () => { - mockListRunners.mockImplementation(async () => [ - ...ec2InstancesRegistered, - { - id: 'i-4-still-booting', - launchTime: moment(new Date()) - .subtract(MINIMUM_TIME_RUNNING - 3, 'minutes') - .toDate(), - type: 'Org', - owner: ORG, - }, - { - id: 'i-5-orphan', - launchTime: moment(new Date()) - .subtract(MINIMUM_TIME_RUNNING + 3, 'minutes') - .toDate(), - type: 'Org', - owner: ORG, - }, - ]); - // 2 idle + 1 booting = 3, top up with 2 to match a pool of 5 - await adjust({ poolSize: 5, type: 'ec2' }); - expect(createRunners).toHaveBeenCalled(); - // Access the numberOfRunners without assuming a specific position - // Just test that the function was called - expect(createRunners).toHaveBeenCalled(); - // With TypeScript we can't directly access mock.calls, so we'll just verify the function was called - // The number of runners should be correct, but we can't type-check this easily + expect(poolProvider.listRunners).not.toHaveBeenCalled(); }); - it('Should not top up if pool size is reached including a booting instance.', async () => { - mockListRunners.mockImplementation(async () => [ - ...ec2InstancesRegistered, - { - id: 'i-4-still-booting', - launchTime: moment(new Date()) - .subtract(MINIMUM_TIME_RUNNING - 3, 'minutes') - .toDate(), - type: 'Org', - owner: ORG, - }, - { - id: 'i-5-orphan', - launchTime: moment(new Date()) - .subtract(MINIMUM_TIME_RUNNING + 3, 'minutes') - .toDate(), - type: 'Org', - owner: ORG, - }, - ]); + it('does not top up when the requested size is already available', async () => { + await adjust({ poolSize: 1 }); - await adjust({ poolSize: 2, type: 'ec2' }); - expect(createRunners).not.toHaveBeenCalled(); + expect(poolProvider.createRunners).not.toHaveBeenCalled(); }); }); describe('With GHES', () => { beforeEach(() => { - (getGitHubEnterpriseApiUrl as ReturnType).mockReturnValue({ + mockedGetGitHubEnterpriseApiUrl.mockReturnValue({ ghesApiUrl: 'https://api.github.enterprise.something', ghesBaseUrl: 'https://github.enterprise.something', }); }); - it('Top up if the pool size is set to 5', async () => { - await adjust({ poolSize: 5, type: 'ec2' }); - // 2 idle, top up with 3 to match a pool of 5 - expect(createRunners).toHaveBeenCalledWith( - expect.anything(), - expect.anything(), - 3, - expect.anything(), - expect.anything(), - 'pool-lambda', + it('passes the enterprise base URL when creating runners', async () => { + await adjust({ poolSize: 5 }); + + expect(poolProvider.createRunners).toHaveBeenCalledWith( + expect.objectContaining({ + githubRunnerConfig: expect.objectContaining({ + ghesBaseUrl: 'https://github.enterprise.something', + }), + numberOfRunners: 3, + }), ); }); }); - describe('With Github Data Residency', () => { + describe('With GitHub Data Residency', () => { beforeEach(() => { - (getGitHubEnterpriseApiUrl as ReturnType).mockReturnValue({ + mockedGetGitHubEnterpriseApiUrl.mockReturnValue({ ghesApiUrl: 'https://api.companyname.ghe.com', ghesBaseUrl: 'https://companyname.ghe.com', }); }); - it('Top up if the pool size is set to 5', async () => { - await adjust({ poolSize: 5, type: 'ec2' }); - // 2 idle, top up with 3 to match a pool of 5 - expect(createRunners).toHaveBeenCalledWith( - expect.anything(), - expect.anything(), - 3, - expect.anything(), - expect.anything(), - 'pool-lambda', + it('passes the data-residency base URL when creating runners', async () => { + await adjust({ poolSize: 5 }); + + expect(poolProvider.createRunners).toHaveBeenCalledWith( + expect.objectContaining({ + githubRunnerConfig: expect.objectContaining({ + ghesBaseUrl: 'https://companyname.ghe.com', + }), + numberOfRunners: 3, + }), ); }); }); @@ -364,153 +217,98 @@ describe('Test simple pool.', () => { process.env.RUNNER_NAME_PREFIX = 'runner-prefix_'; }); - it('Should top up with fewer runners when there are idle prefixed runners', async () => { - // Add prefixed runners to github - mockOctokit.paginate.mockImplementation(async () => [ + it('removes the prefix from runner names before passing statuses to the provider', async () => { + githubClient.paginate.mockResolvedValue([ ...githubRunnersRegistered, { id: 5, - name: 'runner-prefix_i-5-idle', - os: 'linux', - status: 'online', - busy: false, - labels: [], - }, - { - id: 6, - name: 'runner-prefix_i-6-idle', + name: 'runner-prefix_runner-5', os: 'linux', status: 'online', busy: false, labels: [], }, ]); + poolProvider.countAvailableRunners.mockReturnValue(4); - // Add instances in ec2 - mockListRunners.mockImplementation(async () => [ - ...ec2InstancesRegistered, - { - id: 'i-5-idle', - launchTime: new Date(), - type: 'Org', - owner: ORG, - }, - { - id: 'i-6-idle', - launchTime: new Date(), - type: 'Org', - owner: ORG, - }, - ]); + await adjust({ poolSize: 5 }); - await adjust({ poolSize: 5, type: 'ec2' }); - // 2 idle, 2 prefixed idle top up with 1 to match a pool of 5 - expect(createRunners).toHaveBeenCalledWith( - expect.anything(), - expect.anything(), - 1, - expect.anything(), - expect.anything(), - 'pool-lambda', - ); + expect(poolProvider.countAvailableRunners).toHaveBeenCalledWith(providerRunners, expect.any(Map), false); + const runnerStatuses = poolProvider.countAvailableRunners.mock.calls[0][1]; + expect(runnerStatuses.get('runner-5')).toEqual({ busy: false, status: 'online' }); + expect(runnerStatuses.has('runner-prefix_runner-5')).toBe(false); + expect(poolProvider.createRunners).toHaveBeenCalledWith(expect.objectContaining({ numberOfRunners: 1 })); }); }); describe('Respecting runners_maximum_count', () => { beforeEach(() => { - (getGitHubEnterpriseApiUrl as ReturnType).mockReturnValue({ + mockedGetGitHubEnterpriseApiUrl.mockReturnValue({ ghesApiUrl: '', ghesBaseUrl: '', }); }); - it('Should not top up when the total number of running runners is at the maximum.', async () => { - // 4 running runners (2 idle, 1 busy, 1 offline) already meet the maximum, so a large pool size - // must not create more. This is the over-provisioning case from issue #5186. + it('does not top up when the total number of runners is at the maximum', async () => { process.env.RUNNERS_MAXIMUM_COUNT = '4'; - await adjust({ poolSize: 10, type: 'ec2' }); - expect(createRunners).not.toHaveBeenCalled(); + await adjust({ poolSize: 10 }); + + expect(poolProvider.createRunners).not.toHaveBeenCalled(); }); - it('Should not top up when the total number of running runners exceeds the maximum.', async () => { + it('does not top up when the total number of runners exceeds the maximum', async () => { process.env.RUNNERS_MAXIMUM_COUNT = '3'; - await adjust({ poolSize: 10, type: 'ec2' }); - expect(createRunners).not.toHaveBeenCalled(); + await adjust({ poolSize: 10 }); + + expect(poolProvider.createRunners).not.toHaveBeenCalled(); }); - it('Should clamp the top-up to the remaining headroom under the maximum.', async () => { - // 4 running runners with a maximum of 6 leaves headroom for 2, even though the pool of 10 and the - // 2 idle runners would otherwise request a top-up of 8. + it('clamps the top-up to the remaining headroom under the maximum', async () => { process.env.RUNNERS_MAXIMUM_COUNT = '6'; - await adjust({ poolSize: 10, type: 'ec2' }); - expect(createRunners).toHaveBeenCalledWith( - expect.anything(), - expect.anything(), - 2, - expect.anything(), - expect.anything(), - 'pool-lambda', - ); + await adjust({ poolSize: 10 }); + + expect(poolProvider.createRunners).toHaveBeenCalledWith(expect.objectContaining({ numberOfRunners: 2 })); }); - it('Should top up against the pool size when below the maximum headroom.', async () => { - // Headroom (6 - 4 = 2) is larger than the pool demand (5 - 2 idle = 3 would exceed it, so use a - // pool that stays within headroom): pool of 3 with 2 idle requests 1, which is under the cap. + it('tops up against the pool size when below the maximum headroom', async () => { process.env.RUNNERS_MAXIMUM_COUNT = '6'; - await adjust({ poolSize: 3, type: 'ec2' }); - expect(createRunners).toHaveBeenCalledWith( - expect.anything(), - expect.anything(), - 1, - expect.anything(), - expect.anything(), - 'pool-lambda', - ); + await adjust({ poolSize: 3 }); + + expect(poolProvider.createRunners).toHaveBeenCalledWith(expect.objectContaining({ numberOfRunners: 1 })); }); - it('Should ignore the maximum when set to -1 (unlimited).', async () => { + it('ignores the maximum when set to -1', async () => { process.env.RUNNERS_MAXIMUM_COUNT = '-1'; - // 2 idle of 4 running, pool of 10 tops up with 8 regardless of how many are already running. - await adjust({ poolSize: 10, type: 'ec2' }); - expect(createRunners).toHaveBeenCalledWith( - expect.anything(), - expect.anything(), - 8, - expect.anything(), - expect.anything(), - 'pool-lambda', - ); + await adjust({ poolSize: 10 }); + + expect(poolProvider.createRunners).toHaveBeenCalledWith(expect.objectContaining({ numberOfRunners: 8 })); }); }); describe('With INCLUDE_BUSY_RUNNERS enabled', () => { beforeEach(() => { process.env.INCLUDE_BUSY_RUNNERS = 'true'; + poolProvider.countAvailableRunners.mockReturnValue(3); }); - it('Should not top up when pool size matches runners including busy online runners.', async () => { - // Without INCLUDE_BUSY_RUNNERS: 2 in pool (i-1-idle, i-4-idle-older). With it: 3 (adds i-2-busy). - await adjust({ poolSize: 3, type: 'ec2' }); - expect(createRunners).not.toHaveBeenCalled(); + it('does not top up when the provider reports that the requested size is available', async () => { + await adjust({ poolSize: 3 }); + + expect(poolProvider.countAvailableRunners).toHaveBeenCalledWith(providerRunners, expect.any(Map), true); + expect(poolProvider.createRunners).not.toHaveBeenCalled(); }); - it('Should top up by two runners when pool size is 5 and busy runners count toward the pool.', async () => { - await adjust({ poolSize: 5, type: 'ec2' }); - // 3 in pool (idle, busy, older idle); need 2 more - expect(createRunners).toHaveBeenCalledWith( - expect.anything(), - expect.anything(), - 2, - expect.anything(), - expect.anything(), - 'pool-lambda', - ); + it('tops up from the provider count that includes busy runners', async () => { + await adjust({ poolSize: 5 }); + + expect(poolProvider.countAvailableRunners).toHaveBeenCalledWith(providerRunners, expect.any(Map), true); + expect(poolProvider.createRunners).toHaveBeenCalledWith(expect.objectContaining({ numberOfRunners: 2 })); }); }); describe('Multi-app round-robin', () => { beforeEach(() => { - (getGitHubEnterpriseApiUrl as ReturnType).mockReturnValue({ + mockedGetGitHubEnterpriseApiUrl.mockReturnValue({ ghesApiUrl: '', ghesBaseUrl: '', }); @@ -527,11 +325,7 @@ describe('Test simple pool.', () => { await adjust({ poolSize: 3 }); - expect(mockedInstallationAuth).toHaveBeenCalledWith( - expect.any(Number), - expect.any(String), - 1, // appIndex must match the one from createGithubAppAuth - ); + expect(mockedInstallationAuth).toHaveBeenCalledWith(expect.any(Number), expect.any(String), 1); }); it('looks up installationId using the selected app JWT', async () => { @@ -545,11 +339,10 @@ describe('Test simple pool.', () => { await adjust({ poolSize: 3 }); - // Should look up installationId via the API - expect(mockOctokit.apps.getOrgInstallation).toHaveBeenCalledWith({ org: ORG }); + expect(githubClient.apps.getOrgInstallation).toHaveBeenCalledWith({ org: ORG }); }); - it('passes appIndex to createRunners so rate-limit metrics are attributed to the correct app', async () => { + it('passes appIndex to the provider so rate-limit metrics use the selected app', async () => { mockedAppAuth.mockResolvedValue({ type: 'app', token: 'token', @@ -560,14 +353,10 @@ describe('Test simple pool.', () => { await adjust({ poolSize: 3 }); - expect(createRunners).toHaveBeenCalledWith( - // appIndex must match the one returned by createGithubAppAuth - expect.objectContaining({ appIndex: 2 }), - expect.anything(), - expect.any(Number), - expect.anything(), - expect.anything(), - 'pool-lambda', + expect(poolProvider.createRunners).toHaveBeenCalledWith( + expect.objectContaining({ + githubRunnerConfig: expect.objectContaining({ appIndex: 2 }), + }), ); }); }); diff --git a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts index 148da04d93..5667cf9bc0 100644 --- a/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts +++ b/lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts @@ -1,13 +1,24 @@ -import type { RunnerInfo } from '../../../../core'; -import { bootTimeExceeded } from '../runners'; -import { calculateEc2PoolSize } from './pool'; +import type { Octokit } from '@octokit/rest'; +import type { CreateGitHubRunnerConfig, CreateStartRunnerConfig, RunnerInfo } from '../../../../core'; +import { bootTimeExceeded, listEC2Runners } from '../runners'; +import { calculateEc2PoolSize, createEc2PoolProvider } from './pool'; +import { createRunners, type Ec2ProviderConfig, loadEc2ProviderConfig } from './runner-config'; import { beforeEach, describe, expect, it, vi } from 'vitest'; vi.mock('../runners', () => ({ bootTimeExceeded: vi.fn(), + listEC2Runners: vi.fn(), +})); + +vi.mock('./runner-config', () => ({ + createRunners: vi.fn(), + loadEc2ProviderConfig: vi.fn(), })); const mockBootTimeExceeded = vi.mocked(bootTimeExceeded); +const mockListRunners = vi.mocked(listEC2Runners); +const mockCreateRunners = vi.mocked(createRunners); +const mockLoadProviderConfig = vi.mocked(loadEc2ProviderConfig); describe('calculateEc2PoolSize', () => { beforeEach(() => { @@ -36,6 +47,14 @@ describe('calculateEc2PoolSize', () => { expect(mockBootTimeExceeded).not.toHaveBeenCalled(); }); + it('counts registered busy runners when busy runners are included', () => { + const runners: RunnerInfo[] = [{ id: 'i-busy', owner: 'owner', type: 'Org' }]; + const runnerStatus = new Map([['i-busy', { busy: true, status: 'online' }]]); + + expect(calculateEc2PoolSize(runners, runnerStatus, true)).toBe(1); + expect(mockBootTimeExceeded).not.toHaveBeenCalled(); + }); + it('counts unregistered runners that are still booting', () => { const runners: RunnerInfo[] = [{ id: 'i-booting', owner: 'owner', type: 'Org' }]; mockBootTimeExceeded.mockReturnValue(false); @@ -50,3 +69,84 @@ describe('calculateEc2PoolSize', () => { expect(calculateEc2PoolSize(runners, new Map())).toBe(0); }); }); + +describe('createEc2PoolProvider', () => { + const createStartRunnerConfig = vi.fn(); + const githubInstallationClient = {} as Octokit; + const githubRunnerConfig: CreateGitHubRunnerConfig = { + ephemeral: true, + enableJitConfig: true, + runnerLabels: 'self-hosted', + runnerGroup: 'default', + runnerNamePrefix: '', + runnerOwner: 'owner', + runnerType: 'Org', + disableAutoUpdate: false, + ssmTokenPath: '/runners/tokens', + ssmConfigPath: '/runners/config', + ssmParameterStoreTags: [], + }; + const providerConfig: Ec2ProviderConfig = { + environment: 'test-environment', + subnets: ['subnet-123'], + launchTemplateName: 'runner-template', + ec2instanceCriteria: { + instanceTypes: ['m5.large'], + targetCapacityType: 'spot', + instanceAllocationStrategy: 'lowest-price', + }, + tracingEnabled: false, + onDemandFailoverOnError: [], + scaleErrors: [], + }; + + beforeEach(() => { + vi.clearAllMocks(); + mockLoadProviderConfig.mockReturnValue(providerConfig); + }); + + it('lists only running instances managed for the requested pool', async () => { + const runners: RunnerInfo[] = [{ id: 'i-running', owner: 'owner', type: 'Org' }]; + mockListRunners.mockResolvedValue(runners); + const provider = createEc2PoolProvider(createStartRunnerConfig); + + await expect( + provider.listRunners({ + environment: 'test-environment', + runnerOwner: 'owner', + runnerType: 'Org', + }), + ).resolves.toBe(runners); + expect(mockListRunners).toHaveBeenCalledWith({ + environment: 'test-environment', + runnerOwner: 'owner', + runnerType: 'Org', + statuses: ['running'], + }); + }); + + it('creates pool runners with the pool source and returns their instance IDs', async () => { + mockCreateRunners.mockResolvedValue({ + instances: ['i-created'], + retryableErrorCount: 0, + nonRetryableErrorCount: 0, + }); + const provider = createEc2PoolProvider(createStartRunnerConfig); + + await expect( + provider.createRunners({ + githubRunnerConfig, + numberOfRunners: 1, + githubInstallationClient, + }), + ).resolves.toEqual(['i-created']); + expect(mockCreateRunners).toHaveBeenCalledWith( + githubRunnerConfig, + providerConfig, + 1, + githubInstallationClient, + createStartRunnerConfig, + 'pool-lambda', + ); + }); +});