diff --git a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts index f2f71c278d9..74a2c97dd75 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts @@ -13,12 +13,12 @@ import { } from 'aws-core-vscode/codewhisperer' import { assertTelemetry, - captureEventOnce, getTestWindow, SeverityLevel, createBuilderIdProfile, createSsoProfile, createTestAuth, + captureEventNTimes, } from 'aws-core-vscode/test' import { Auth, Connection, isAnySsoConnection, isBuilderIdConnection } from 'aws-core-vscode/auth' import { globals, vscodeComponent } from 'aws-core-vscode/shared' @@ -231,15 +231,13 @@ describe('AuthUtil', async function () { assert.strictEqual(auth.activeConnection?.id, authUtil.conn?.id) // Switch to unsupported connection - const cwAuthUpdatedConnection = captureEventOnce(authUtil.secondaryAuth.onDidChangeActiveConnection) + const cwAuthUpdatedConnection = captureEventNTimes(authUtil.secondaryAuth.onDidChangeActiveConnection, 2) await auth.useConnection(unsupportedConn) - // This is triggered when the main Auth connection is switched + // - This is triggered when the main Auth connection is switched + // - This is triggered by registerAuthListener() when it saves the previous active connection as a fallback. await cwAuthUpdatedConnection - // This is triggered by registerAuthListener() when it saves the previous active connection as a fallback. - // TODO in a refactor see if we can simplify multiple multiple triggers on the same event. - await captureEventOnce(authUtil.secondaryAuth.onDidChangeActiveConnection) - // Is using the fallback connection + // TODO in a refactor see if we can simplify multiple multiple triggers on the same event. assert.ok(authUtil.isConnected()) assert.ok(authUtil.isUsingSavedConnection) assert.notStrictEqual(auth.activeConnection?.id, authUtil.conn?.id) diff --git a/packages/core/src/test/testUtil.ts b/packages/core/src/test/testUtil.ts index 138b802ba7b..26b0c473403 100644 --- a/packages/core/src/test/testUtil.ts +++ b/packages/core/src/test/testUtil.ts @@ -565,9 +565,18 @@ export function captureEvent(event: vscode.Event): EventCapturer { * Captures the first value emitted by an event, optionally with a timeout */ export function captureEventOnce(event: vscode.Event, timeout?: number): Promise { + return captureEventNTimes(event, 1, timeout) +} + +export function captureEventNTimes(event: vscode.Event, amount: number, timeout?: number): Promise { return new Promise((resolve, reject) => { const stop = () => reject(new Error('Timed out waiting for event')) - event((data) => resolve(data)) + let count = 0 + event((data) => { + if (++count === amount) { + resolve(data) + } + }) if (timeout !== undefined) { setTimeout(stop, timeout)