From 3eee9b2578ee3b2fbe588e8ade896f7f89d938e5 Mon Sep 17 00:00:00 2001 From: Stuart Pearson <1926002+stuartp44@users.noreply.github.com> Date: Thu, 18 Dec 2025 10:59:45 +0100 Subject: [PATCH 1/2] fix: improve readability of retry message assertions in tests Co-authored-by: Brend Smits --- .../src/scale-runners/scale-up.test.ts | 160 ++++++++++++++++++ .../src/scale-runners/scale-up.ts | 4 +- 2 files changed, 163 insertions(+), 1 deletion(-) diff --git a/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts b/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts index b876d31d50..9549e7cb2d 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts @@ -10,6 +10,7 @@ import { createRunner, listEC2Runners } from './../aws/runners'; import { RunnerInputParameters } from './../aws/runners.d'; import * as scaleUpModule from './scale-up'; import { getParameter } from '@aws-github-runner/aws-ssm-util'; +import { publishRetryMessage } from './job-retry'; import { describe, it, expect, beforeEach, vi } from 'vitest'; import type { Octokit } from '@octokit/rest'; @@ -33,6 +34,7 @@ const mockCreateRunner = vi.mocked(createRunner); const mockListRunners = vi.mocked(listEC2Runners); const mockSSMClient = mockClient(SSMClient); const mockSSMgetParameter = vi.mocked(getParameter); +const mockPublishRetryMessage = vi.mocked(publishRetryMessage); vi.mock('@octokit/rest', () => ({ Octokit: vi.fn().mockImplementation(function () { @@ -63,6 +65,11 @@ vi.mock('@aws-github-runner/aws-ssm-util', async () => { }; }); +vi.mock('./job-retry', () => ({ + publishRetryMessage: vi.fn(), + checkAndRetryJob: vi.fn(), +})); + export type RunnerType = 'ephemeral' | 'non-ephemeral'; // for ephemeral and non-ephemeral runners @@ -1667,6 +1674,159 @@ describe('scaleUp with Github Data Residency', () => { }); }); +describe('Retry mechanism tests', () => { + beforeEach(() => { + process.env.ENABLE_ORGANIZATION_RUNNERS = 'true'; + process.env.ENABLE_EPHEMERAL_RUNNERS = 'true'; + process.env.ENABLE_JOB_QUEUED_CHECK = 'true'; + process.env.RUNNERS_MAXIMUM_COUNT = '10'; + expectedRunnerParams = { ...EXPECTED_RUNNER_PARAMS }; + mockSSMClient.reset(); + }); + + const createTestMessages = ( + count: number, + overrides: Partial[] = [], + ): scaleUpModule.ActionRequestMessageSQS[] => { + return Array.from({ length: count }, (_, i) => ({ + ...TEST_DATA_SINGLE, + id: i + 1, + messageId: `message-${i + 1}`, + ...overrides[i], + })); + }; + + it('calls publishRetryMessage for each valid message when job is queued', async () => { + const messages = createTestMessages(3); + + await scaleUpModule.scaleUp(messages); + + expect(mockPublishRetryMessage).toHaveBeenCalledTimes(3); + expect(mockPublishRetryMessage).toHaveBeenNthCalledWith( + 1, + expect.objectContaining({ + id: 1, + messageId: 'message-1', + }), + ); + expect(mockPublishRetryMessage).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ + id: 2, + messageId: 'message-2', + }), + ); + expect(mockPublishRetryMessage).toHaveBeenNthCalledWith( + 3, + expect.objectContaining({ + id: 3, + messageId: 'message-3', + }), + ); + }); + + it('does not call publishRetryMessage when job is not queued', async () => { + mockOctokit.actions.getJobForWorkflowRun.mockImplementation((params) => { + const isQueued = params.job_id === 1; // Only job 1 is queued + return { + data: { + status: isQueued ? 'queued' : 'completed', + }, + }; + }); + + const messages = createTestMessages(3); + + await scaleUpModule.scaleUp(messages); + + // Only message with id 1 should trigger retry + expect(mockPublishRetryMessage).toHaveBeenCalledTimes(1); + expect(mockPublishRetryMessage).toHaveBeenCalledWith( + expect.objectContaining({ + id: 1, + messageId: 'message-1', + }), + ); + }); + + it('calls publishRetryMessage even when maximum runners is reached', async () => { + process.env.RUNNERS_MAXIMUM_COUNT = '0'; // No runners can be created + + const messages = createTestMessages(2); + + await scaleUpModule.scaleUp(messages); + + // publishRetryMessage should still be called even though no runners will be created + expect(mockPublishRetryMessage).toHaveBeenCalledTimes(2); + expect(createRunner).not.toHaveBeenCalled(); + }); + + it('calls publishRetryMessage with correct message structure including retry counter', async () => { + const message = { + ...TEST_DATA_SINGLE, + messageId: 'test-message-id', + retryCounter: 2, + }; + + await scaleUpModule.scaleUp([message]); + + expect(mockPublishRetryMessage).toHaveBeenCalledWith( + expect.objectContaining({ + id: message.id, + messageId: 'test-message-id', + retryCounter: 2, + }), + ); + }); + + it('calls publishRetryMessage when ENABLE_JOB_QUEUED_CHECK is false', async () => { + process.env.ENABLE_JOB_QUEUED_CHECK = 'false'; + + const messages = createTestMessages(2); + + await scaleUpModule.scaleUp(messages); + + // Should always call publishRetryMessage when queue check is disabled + expect(mockPublishRetryMessage).toHaveBeenCalledTimes(2); + expect(mockOctokit.actions.getJobForWorkflowRun).not.toHaveBeenCalled(); + }); + + it('calls publishRetryMessage for each message in a multi-runner scenario', async () => { + const messages = createTestMessages(5); + + await scaleUpModule.scaleUp(messages); + + expect(mockPublishRetryMessage).toHaveBeenCalledTimes(5); + messages.forEach((msg, index) => { + expect(mockPublishRetryMessage).toHaveBeenNthCalledWith( + index + 1, + expect.objectContaining({ + id: msg.id, + messageId: msg.messageId, + }), + ); + }); + }); + + it('calls publishRetryMessage before runner creation', async () => { + const messages = createTestMessages(1); + + const callOrder: string[] = []; + mockPublishRetryMessage.mockImplementation(() => { + callOrder.push('publishRetryMessage'); + return Promise.resolve(); + }); + mockCreateRunner.mockImplementation(async () => { + callOrder.push('createRunner'); + return ['i-12345']; + }); + + await scaleUpModule.scaleUp(messages); + + expect(callOrder).toEqual(['publishRetryMessage', 'createRunner']); + }); +}); + function defaultOctokitMockImpl() { mockOctokit.actions.getJobForWorkflowRun.mockImplementation(() => ({ data: { diff --git a/lambdas/functions/control-plane/src/scale-runners/scale-up.ts b/lambdas/functions/control-plane/src/scale-runners/scale-up.ts index 35df7ea5d7..d340b06010 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-up.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-up.ts @@ -7,6 +7,7 @@ import { createGithubAppAuth, createGithubInstallationAuth, createOctokitClient import { createRunner, listEC2Runners, tag } from './../aws/runners'; import { RunnerInputParameters } from './../aws/runners.d'; import { metricGitHubAppRateLimit } from '../github/rate-limit'; +import { publishRetryMessage } from './job-retry'; const logger = createChildLogger('scale-up'); @@ -356,6 +357,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise Date: Sun, 21 Dec 2025 19:58:28 +0100 Subject: [PATCH 2/2] Update lambdas/functions/control-plane/src/scale-runners/scale-up.ts Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- lambdas/functions/control-plane/src/scale-runners/scale-up.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lambdas/functions/control-plane/src/scale-runners/scale-up.ts b/lambdas/functions/control-plane/src/scale-runners/scale-up.ts index d340b06010..d122d06974 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-up.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-up.ts @@ -357,7 +357,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise