From ebcb987198da5db7023ac6b16f8c86a34e913475 Mon Sep 17 00:00:00 2001
From: Nikolas Komonen <118216176+nkomonen-amazon@users.noreply.github.com>
Date: Tue, 15 Oct 2024 16:36:58 -0400
Subject: [PATCH] fix(test): unreliable "AuthUtil CodeWhisperer uses fallback
 connection" #5788

## Problem:

This test was flaky due to the event emitter and the test expecting 2 of
the same event to come in, but sometimes the second one was not capture.

This looks to be due to the second `captureEventOnce` call needing to do
some initial setup before it could capture an event. And I think there
was a race condition between it being setup in time and the event being
emitted before it could start listening.

## Solution:

Make a new function which does the event capture once and then listens N
amount of times. This is reliable.
---
 .../test/unit/codewhisperer/util/authUtil.test.ts    | 12 +++++-------
 packages/core/src/test/testUtil.ts                   | 11 ++++++++++-
 2 files changed, 15 insertions(+), 8 deletions(-)

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<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)