Skip to content

Commit cf7124c

Browse files
gertjanmaasgertjanmaas
committed
WIP: list runners, create token if needed
1 parent 503543c commit cf7124c

File tree

5 files changed

+230
-411
lines changed

5 files changed

+230
-411
lines changed

‎modules/runners/lambdas/scale-runners/src/scale-runners/handler.test.ts

Lines changed: 77 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
1+
import { mocked } from 'ts-jest/utils';
12
import { ActionRequestMessage, handle } from './handler';
23

34
import { createAppAuth } from '@octokit/auth-app';
45
import { Octokit } from '@octokit/rest';
6+
import { listRunners } from './runners';
57

68
jest.mock('@octokit/auth-app', () => ({
79
createAppAuth: jest.fn().mockImplementation(() => jest.fn().mockImplementation(() => ({ token: 'Blaat' }))),
@@ -10,14 +12,16 @@ const mockOctokit = {
1012
checks: { get: jest.fn() },
1113
actions: {
1214
listRepoWorkflowRuns: jest.fn(),
13-
listSelfHostedRunnersForOrg: jest.fn(),
14-
listSelfHostedRunnersForRepo: jest.fn(),
15+
createRegistrationTokenForOrg: jest.fn(),
16+
createRegistrationTokenForRepo: jest.fn(),
1517
},
1618
};
1719
jest.mock('@octokit/rest', () => ({
1820
Octokit: jest.fn().mockImplementation(() => mockOctokit),
1921
}));
2022

23+
jest.mock('./runners');
24+
2125
const TEST_DATA: ActionRequestMessage = {
2226
id: 1,
2327
eventType: 'check_run',
@@ -32,27 +36,29 @@ describe('handler', () => {
3236
process.env.GITHUB_APP_ID = '1337';
3337
process.env.GITHUB_APP_CLIENT_ID = 'TEST_CLIENT_ID';
3438
process.env.GITHUB_APP_CLIENT_SECRET = 'TEST_CLIENT_SECRET';
39+
process.env.RUNNERS_MAXIMUM_COUNT = '3';
3540
jest.clearAllMocks();
3641
mockOctokit.actions.listRepoWorkflowRuns.mockImplementation(() => ({
3742
data: {
3843
total_count: 1,
3944
},
4045
}));
41-
const mockRunnersReturnValue = {
46+
const mockTokenReturnValue = {
4247
data: {
43-
total_count: 1,
44-
runners: [
45-
{
46-
id: 23,
47-
name: 'Test Runner',
48-
status: 'online',
49-
os: 'linux',
50-
},
51-
],
48+
token: '1234abcd',
5249
},
5350
};
54-
mockOctokit.actions.listSelfHostedRunnersForOrg.mockImplementation(() => mockRunnersReturnValue);
55-
mockOctokit.actions.listSelfHostedRunnersForRepo.mockImplementation(() => mockRunnersReturnValue);
51+
mockOctokit.actions.createRegistrationTokenForOrg.mockImplementation(() => mockTokenReturnValue);
52+
mockOctokit.actions.createRegistrationTokenForRepo.mockImplementation(() => mockTokenReturnValue);
53+
const mockListRunners = mocked(listRunners);
54+
mockListRunners.mockImplementation(async () => [
55+
{
56+
instanceId: 'i-1234',
57+
launchTime: new Date(),
58+
repo: `${TEST_DATA.repositoryOwner}/${TEST_DATA.repositoryName}`,
59+
org: TEST_DATA.repositoryOwner,
60+
},
61+
]);
5662
});
5763

5864
it('ignores non-sqs events', async () => {
@@ -69,30 +75,61 @@ describe('handler', () => {
6975
});
7076
});
7177

72-
// describe('on org level', () => {
73-
// beforeAll(() => {
74-
// process.env.ENABLE_ORGANIZATION_RUNNERS = 'true';
75-
// });
76-
77-
// it('gets the current org level runners', async () => {
78-
// await handle('aws:sqs', TEST_DATA);
79-
// expect(mockOctokit.actions.listSelfHostedRunnersForOrg).toBeCalledWith({
80-
// org: TEST_DATA.repositoryOwner,
81-
// });
82-
// });
83-
// });
84-
85-
// describe('on repo level', () => {
86-
// beforeAll(() => {
87-
// delete process.env.ENABLE_ORGANIZATION_RUNNERS;
88-
// });
89-
90-
// it('gets the current repo level runners', async () => {
91-
// await handle('aws:sqs', TEST_DATA);
92-
// expect(mockOctokit.actions.listSelfHostedRunnersForRepo).toBeCalledWith({
93-
// owner: TEST_DATA.repositoryOwner,
94-
// repo: TEST_DATA.repositoryName,
95-
// });
96-
// });
97-
// });
78+
it('does not list runners when no workflows are queued', async () => {
79+
mockOctokit.actions.listRepoWorkflowRuns.mockImplementation(() => ({
80+
data: { total_count: 0, runners: [] },
81+
}));
82+
await handle('aws:sqs', TEST_DATA);
83+
expect(listRunners).not.toBeCalled();
84+
});
85+
86+
describe('on org level', () => {
87+
beforeAll(() => {
88+
process.env.ENABLE_ORGANIZATION_RUNNERS = 'true';
89+
});
90+
91+
it('gets the current org level runners', async () => {
92+
await handle('aws:sqs', TEST_DATA);
93+
expect(listRunners).toBeCalledWith({ repoName: undefined });
94+
});
95+
96+
it('does not create a token when maximum runners has been reached', async () => {
97+
process.env.RUNNERS_MAXIMUM_COUNT = '1';
98+
await handle('aws:sqs', TEST_DATA);
99+
expect(mockOctokit.actions.createRegistrationTokenForOrg).not.toBeCalled();
100+
});
101+
102+
it('creates a token when maximum runners has not been reached', async () => {
103+
await handle('aws:sqs', TEST_DATA);
104+
expect(mockOctokit.actions.createRegistrationTokenForOrg).toBeCalled();
105+
expect(mockOctokit.actions.createRegistrationTokenForOrg).toBeCalledWith({
106+
org: TEST_DATA.repositoryOwner,
107+
});
108+
});
109+
});
110+
111+
describe('on repo level', () => {
112+
beforeAll(() => {
113+
process.env.ENABLE_ORGANIZATION_RUNNERS = 'false';
114+
});
115+
116+
it('gets the current repo level runners', async () => {
117+
await handle('aws:sqs', TEST_DATA);
118+
expect(listRunners).toBeCalledWith({ repoName: `${TEST_DATA.repositoryOwner}/${TEST_DATA.repositoryName}` });
119+
});
120+
121+
it('does not create a token when maximum runners has been reached', async () => {
122+
process.env.RUNNERS_MAXIMUM_COUNT = '1';
123+
await handle('aws:sqs', TEST_DATA);
124+
expect(mockOctokit.actions.createRegistrationTokenForRepo).not.toBeCalled();
125+
});
126+
127+
it('creates a token when maximum runners has not been reached', async () => {
128+
await handle('aws:sqs', TEST_DATA);
129+
expect(mockOctokit.actions.createRegistrationTokenForRepo).toBeCalledWith({
130+
owner: TEST_DATA.repositoryOwner,
131+
repo: TEST_DATA.repositoryName,
132+
});
133+
});
134+
});
98135
});

‎modules/runners/lambdas/scale-runners/src/scale-runners/handler.ts

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { createAppAuth } from '@octokit/auth-app';
22
import { Octokit } from '@octokit/rest';
33
import { AppAuth } from '@octokit/auth-app/dist-types/types';
4+
import { listRunners } from './runners';
45
import yn from 'yn';
56

67
export interface ActionRequestMessage {
@@ -34,29 +35,42 @@ async function createInstallationClient(githubAppAuth: AppAuth): Promise<Octokit
3435
export const handle = async (eventSource: string, payload: ActionRequestMessage): Promise<void> => {
3536
if (eventSource !== 'aws:sqs') throw Error('Cannot handle non-SQS events!');
3637
const enableOrgLevel = yn(process.env.ENABLE_ORGANIZATION_RUNNERS);
38+
const maximumRunners = parseInt(process.env.RUNNERS_MAXIMUM_COUNT || '3');
3739
const githubAppAuth = createGithubAppAuth(payload.installationId);
3840
const githubInstallationClient = await createInstallationClient(githubAppAuth);
3941
const queuedWorkflows = await githubInstallationClient.actions.listRepoWorkflowRuns({
4042
owner: payload.repositoryOwner,
4143
repo: payload.repositoryName,
42-
// @ts-ignore (typing is incorrect)
44+
// @ts-ignore (typing of the 'status' field is incorrect)
4345
status: 'queued',
4446
});
4547
console.info(
4648
`Repo ${payload.repositoryOwner}/${payload.repositoryName} has ${queuedWorkflows.data.total_count} queued workflow runs`,
4749
);
4850

4951
if (queuedWorkflows.data.total_count > 0) {
50-
// console.log(enableOrgLevel);
51-
// const currentRunners = enableOrgLevel
52-
// ? await githubInstallationClient.actions.listSelfHostedRunnersForOrg({
53-
// org: payload.repositoryOwner,
54-
// })
55-
// : await githubInstallationClient.actions.listSelfHostedRunnersForRepo({
56-
// owner: payload.repositoryOwner,
57-
// repo: payload.repositoryName,
58-
// });
59-
// // const currentOnlineRunners = currentRunners.data.runners.filter((r) => r.status === 'online');
60-
// // if (currentOnlineRunners.length > 0)
52+
const currentRunners = await listRunners({
53+
repoName: enableOrgLevel ? undefined : `${payload.repositoryOwner}/${payload.repositoryName}`,
54+
});
55+
console.info(
56+
`${
57+
enableOrgLevel
58+
? `Organization ${payload.repositoryOwner}`
59+
: `Repo ${payload.repositoryOwner}/${payload.repositoryName}`
60+
} has ${currentRunners.length}/${maximumRunners} runners`,
61+
);
62+
console.log(currentRunners.length);
63+
console.log(maximumRunners);
64+
if (currentRunners.length < maximumRunners) {
65+
// create token
66+
const registrationToken = enableOrgLevel
67+
? await githubInstallationClient.actions.createRegistrationTokenForOrg({ org: payload.repositoryOwner })
68+
: await githubInstallationClient.actions.createRegistrationTokenForRepo({
69+
owner: payload.repositoryOwner,
70+
repo: payload.repositoryName,
71+
});
72+
const token = registrationToken.data.token;
73+
// create runner
74+
}
6175
}
6276
};

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

Lines changed: 86 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,98 @@
1-
import { listRunners } from './runners';
2-
import { handle } from './handler';
1+
import { listRunners, RunnerInfo } from './runners';
32
import { EC2 } from 'aws-sdk';
43

5-
jest.mock('./handler');
64
const mockEC2 = { describeInstances: jest.fn() };
75
jest.mock('aws-sdk', () => ({
86
EC2: jest.fn().mockImplementation(() => mockEC2),
97
}));
108

119
describe('list instances', () => {
12-
beforeAll(() => {
10+
const mockDescribeInstances = { promise: jest.fn() };
11+
beforeEach(() => {
1312
jest.clearAllMocks();
13+
mockEC2.describeInstances.mockImplementation(() => mockDescribeInstances);
14+
const mockRunningInstances: AWS.EC2.DescribeInstancesResult = {
15+
Reservations: [
16+
{
17+
Instances: [
18+
{
19+
LaunchTime: new Date('2020-10-10T14:48:00.000+09:00'),
20+
InstanceId: 'i-1234',
21+
Tags: [
22+
{ Key: 'Repo', Value: 'CoderToCat/hello-world' },
23+
{ Key: 'Org', Value: 'CoderToCat' },
24+
{ Key: 'Application', Value: 'github-action-runner' },
25+
],
26+
},
27+
{
28+
LaunchTime: new Date('2020-10-11T14:48:00.000+09:00'),
29+
InstanceId: 'i-5678',
30+
Tags: [
31+
{ Key: 'Repo', Value: 'SomeAwesomeCoder/some-amazing-library' },
32+
{ Key: 'Org', Value: 'SomeAwesomeCoder' },
33+
{ Key: 'Application', Value: 'github-action-runner' },
34+
],
35+
},
36+
],
37+
},
38+
],
39+
};
40+
mockDescribeInstances.promise.mockReturnValue(mockRunningInstances);
1441
});
15-
it('returns a list of instances', () => {
16-
listRunners();
42+
43+
it('returns a list of instances', async () => {
44+
const resp = await listRunners();
45+
expect(resp.length).toBe(2);
46+
expect(resp).toContainEqual({
47+
instanceId: 'i-1234',
48+
launchTime: new Date('2020-10-10T14:48:00.000+09:00'),
49+
repo: 'CoderToCat/hello-world',
50+
org: 'CoderToCat',
51+
});
52+
expect(resp).toContainEqual({
53+
instanceId: 'i-5678',
54+
launchTime: new Date('2020-10-11T14:48:00.000+09:00'),
55+
repo: 'SomeAwesomeCoder/some-amazing-library',
56+
org: 'SomeAwesomeCoder',
57+
});
58+
});
59+
60+
it('calls EC2 describe instances', async () => {
61+
await listRunners();
62+
expect(mockEC2.describeInstances).toBeCalled();
63+
});
64+
65+
it('filters instances on repo name', async () => {
66+
await listRunners({ repoName: 'SomeAwesomeCoder/some-amazing-library' });
67+
expect(mockEC2.describeInstances).toBeCalledWith({
68+
Filters: [
69+
{ Name: 'tag:Application', Values: ['github-action-runner'] },
70+
{ Name: 'instance-state-name', Values: ['running', 'pending'] },
71+
{ Name: 'tag:Repo', Values: ['SomeAwesomeCoder/some-amazing-library'] },
72+
],
73+
});
74+
});
75+
76+
it('filters instances on org name', async () => {
77+
await listRunners({ orgName: 'SomeAwesomeCoder' });
78+
expect(mockEC2.describeInstances).toBeCalledWith({
79+
Filters: [
80+
{ Name: 'tag:Application', Values: ['github-action-runner'] },
81+
{ Name: 'instance-state-name', Values: ['running', 'pending'] },
82+
{ Name: 'tag:Org', Values: ['SomeAwesomeCoder'] },
83+
],
84+
});
85+
});
86+
87+
it('filters instances on both org name and repo name', async () => {
88+
await listRunners({ orgName: 'SomeAwesomeCoder', repoName: 'SomeAwesomeCoder/some-amazing-library' });
89+
expect(mockEC2.describeInstances).toBeCalledWith({
90+
Filters: [
91+
{ Name: 'tag:Application', Values: ['github-action-runner'] },
92+
{ Name: 'instance-state-name', Values: ['running', 'pending'] },
93+
{ Name: 'tag:Repo', Values: ['SomeAwesomeCoder/some-amazing-library'] },
94+
{ Name: 'tag:Org', Values: ['SomeAwesomeCoder'] },
95+
],
96+
});
1797
});
1898
});

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

Lines changed: 34 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,26 +2,45 @@ import { EC2 } from 'aws-sdk';
22

33
export interface RunnerInfo {
44
instanceId: string;
5-
launchTime: Date;
6-
repo: string;
7-
org: string;
5+
launchTime: Date | undefined;
6+
repo: string | undefined;
7+
org: string | undefined;
88
}
99

10-
const ec2 = new EC2();
11-
export async function listRunners(
12-
repoName: string | undefined = undefined,
13-
orgName: string | undefined = undefined,
14-
): Promise<RunnerInfo[]> {
15-
let filters = [
10+
export interface ListRunnerFilters {
11+
repoName?: string;
12+
orgName?: string;
13+
}
14+
15+
export async function listRunners(filters: ListRunnerFilters | undefined = undefined): Promise<RunnerInfo[]> {
16+
const ec2 = new EC2();
17+
let ec2Filters = [
1618
{ Name: 'tag:Application', Values: ['github-action-runner'] },
1719
{ Name: 'instance-state-name', Values: ['running', 'pending'] },
1820
];
19-
if (repoName !== undefined) {
20-
filters.push({ Name: 'tag:Repo', Values: [repoName] });
21+
if (filters) {
22+
if (filters.repoName !== undefined) {
23+
ec2Filters.push({ Name: 'tag:Repo', Values: [filters.repoName] });
24+
}
25+
if (filters.orgName !== undefined) {
26+
ec2Filters.push({ Name: 'tag:Org', Values: [filters.orgName] });
27+
}
2128
}
22-
if (orgName !== undefined) {
23-
filters.push({ Name: 'tag:Org', Values: [orgName] });
29+
const runningInstances = await ec2.describeInstances({ Filters: ec2Filters }).promise();
30+
const runners: RunnerInfo[] = [];
31+
if (runningInstances.Reservations) {
32+
for (const r of runningInstances.Reservations) {
33+
if (r.Instances) {
34+
for (const i of r.Instances) {
35+
runners.push({
36+
instanceId: i.InstanceId as string,
37+
launchTime: i.LaunchTime,
38+
repo: i.Tags?.find((e) => e.Key === 'Repo')?.Value,
39+
org: i.Tags?.find((e) => e.Key === 'Org')?.Value,
40+
});
41+
}
42+
}
43+
}
2444
}
25-
const runningInstances = await ec2.describeInstances({ Filters: filters }).promise();
26-
return [{ instanceId: 'i-123', launchTime: new Date(), repo: 'bla', org: 'bla' }];
45+
return runners;
2746
}

0 commit comments

Comments
 (0)