diff --git a/packages/lambda-durable-functions-sdk-js/src/utils/wait-before-continue/wait-before-continue.test.ts b/packages/lambda-durable-functions-sdk-js/src/utils/wait-before-continue/wait-before-continue.test.ts index 074a784f..485f64e7 100644 --- a/packages/lambda-durable-functions-sdk-js/src/utils/wait-before-continue/wait-before-continue.test.ts +++ b/packages/lambda-durable-functions-sdk-js/src/utils/wait-before-continue/wait-before-continue.test.ts @@ -5,12 +5,20 @@ import { OperationStatus, Operation } from '@amzn/dex-internal-sdk'; describe('waitBeforeContinue', () => { let mockContext: jest.Mocked; let mockHasRunningOperations: jest.Mock; + let timers: NodeJS.Timeout[] = []; beforeEach(() => { mockContext = { getStepData: jest.fn(), } as any; mockHasRunningOperations = jest.fn(); + timers = []; + }); + + afterEach(() => { + // Clean up any remaining timers + timers.forEach(timer => clearTimeout(timer)); + timers = []; }); test('should resolve when operations complete', async () => { @@ -28,9 +36,10 @@ describe('waitBeforeContinue', () => { }); // Complete operations after 50ms - setTimeout(() => { + const timer = setTimeout(() => { operationsRunning = false; }, 50); + timers.push(timer); const result = await resultPromise; expect(result.reason).toBe('operations'); @@ -68,9 +77,10 @@ describe('waitBeforeContinue', () => { }); // Change status after 50ms - setTimeout(() => { + const timer = setTimeout(() => { stepStatus = OperationStatus.SUCCEEDED; }, 50); + timers.push(timer); const result = await resultPromise; expect(result.reason).toBe('status'); diff --git a/packages/lambda-durable-functions-sdk-js/src/utils/wait-before-continue/wait-before-continue.ts b/packages/lambda-durable-functions-sdk-js/src/utils/wait-before-continue/wait-before-continue.ts index 34b1c1dd..7ab03762 100644 --- a/packages/lambda-durable-functions-sdk-js/src/utils/wait-before-continue/wait-before-continue.ts +++ b/packages/lambda-durable-functions-sdk-js/src/utils/wait-before-continue/wait-before-continue.ts @@ -54,13 +54,20 @@ export async function waitBeforeContinue( } = options; const promises: Promise[] = []; + const timers: NodeJS.Timeout[] = []; + + // Cleanup function to clear all timers + const cleanup = () => { + timers.forEach(timer => clearTimeout(timer)); + }; // Timer promise - resolves when scheduled time is reached if (checkTimer && scheduledTimestamp) { const timerPromise = new Promise(resolve => { const timeLeft = Number(scheduledTimestamp) - Date.now(); if (timeLeft > 0) { - setTimeout(() => resolve({ reason: 'timer', timerExpired: true }), timeLeft); + const timer = setTimeout(() => resolve({ reason: 'timer', timerExpired: true }), timeLeft); + timers.push(timer); } else { resolve({ reason: 'timer', timerExpired: true }); } @@ -75,7 +82,8 @@ export async function waitBeforeContinue( if (!hasRunningOperations()) { resolve({ reason: 'operations' }); } else { - setTimeout(checkOperations, pollingInterval); + const timer = setTimeout(checkOperations, pollingInterval); + timers.push(timer); } }; checkOperations(); @@ -92,7 +100,8 @@ export async function waitBeforeContinue( if (originalStatus !== currentStatus) { resolve({ reason: 'status' }); } else { - setTimeout(checkStepStatus, pollingInterval); + const timer = setTimeout(checkStepStatus, pollingInterval); + timers.push(timer); } }; checkStepStatus(); @@ -105,8 +114,9 @@ export async function waitBeforeContinue( return { reason: 'timeout' }; } - // Wait for any condition to be met + // Wait for any condition to be met, then cleanup timers const result = await Promise.race(promises); + cleanup(); // If timer expired, force checkpoint to get fresh data from API if (result.reason === 'timer' && result.timerExpired && checkpoint) {