fix(test): unreliable test "AuthUtil CodeWhisperer uses fallback connection"

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)

packages/core/src/test/testUtil.ts

@@ -565,9 +565,18 @@ export function captureEvent<T>(event: vscode.Event<T>): EventCapturer<T> {
  * Captures the first value emitted by an event, optionally with a timeout
  */
 export function captureEventOnce<T>(event: vscode.Event<T>, timeout?: number): Promise<T> {
+    return captureEventNTimes(event, 1, timeout)
+}
+
+export function captureEventNTimes<T>(event: vscode.Event<T>, amount: number, timeout?: number): Promise<T> {
     return new Promise<T>((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)