Fix flakey test (#1178)

1 files changed, 9 insertions, 2 deletions

diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts
index cd06a11a..47bd45b3 100644
--- a/packages/core/src/code_assist/oauth2.test.ts
+++ b/packages/core/src/code_assist/oauth2.test.ts
@@ -67,11 +67,18 @@ describe('oauth2', () => {
       typeof http.IncomingMessage,
       typeof http.ServerResponse
     >;
+
+    let serverListeningCallback: (value: unknown) => void;
+    const serverListeningPromise = new Promise(
+      (resolve) => (serverListeningCallback = resolve),
+    );
+
     const mockHttpServer = {
       listen: vi.fn((port: number, callback?: () => void) => {
         if (callback) {
           callback();
         }
+        serverListeningCallback(undefined);
       }),
       close: vi.fn((callback?: () => void) => {
         if (callback) {
@@ -91,8 +98,8 @@ describe('oauth2', () => {
 
     const clientPromise = getOauthClient();
 
-    // Wait for the server to be created
-    await new Promise((resolve) => setTimeout(resolve, 0));
+    // wait for server to start listening.
+    await serverListeningPromise;
 
     const mockReq = {
       url: `/oauth2callback?code=${mockCode}&state=${mockState}`,