created waitUntilMcpClientIsSet and use test helper

Author: jeroenvervaeke

package-lock.json

@@ -15,6 +15,7 @@ import { MCPConnectionManager } from "../../src/common/connectionManager.js";
 import { DeviceId } from "../../src/helpers/deviceId.js";
 import { connectionErrorHandler } from "../../src/common/connectionErrorHandler.js";
 import { Keychain } from "../../src/common/keychain.js";
+import type { Client as AtlasLocalClient } from "@mongodb-js-preview/atlas-local";
 
 interface ParameterInfo {
     name: string;
@@ -345,6 +346,37 @@ export function waitUntil<T extends ConnectionState>(
     });
 }
 
+export function waitUntilMcpClientIsSet(
+    mcpServer: Server,
+    signal: AbortSignal,
+    timeout: number = 5000
+): Promise<AtlasLocalClient> {
+    let ts: NodeJS.Timeout | undefined;
+
+    const timeoutSignal = AbortSignal.timeout(timeout);
+    const combinedSignal = AbortSignal.any([signal, timeoutSignal]);
+
+    return new Promise<AtlasLocalClient>((resolve, reject) => {
+        ts = setInterval(() => {
+            if (combinedSignal.aborted) {
+                return reject(new Error(`Aborted: ${combinedSignal.reason}`));
+            }
+
+            // wait until session.client != undefined
+            // do not wait more than 1 second, should take a few milliseconds at most
+            // try every 50ms to see if the client is set, if it's not set after 1 second, throw an error
+            const client = mcpServer.session.atlasLocalClient;
+            if (client) {
+                return resolve(client);
+            }
+        }, 100);
+    }).finally(() => {
+        if (ts !== undefined) {
+            clearInterval(ts);
+        }
+    });
+}
+
 export function getDataFromUntrustedContent(content: string): string {
     const regex = /^[ \t]*<untrusted-user-data-[0-9a-f\\-]*>(?<data>.*)^[ \t]*<\/untrusted-user-data-[0-9a-f\\-]*>/gms;
     const match = regex.exec(content);

tests/integration/helpers.ts

@@ -4,6 +4,7 @@ import {
     expectDefined,
     getResponseElements,
     setupIntegrationTest,
+    waitUntilMcpClientIsSet,
 } from "../../helpers.js";
 import { describe, expect, it } from "vitest";
 
@@ -17,22 +18,26 @@ describe("atlas-local-list-deployments", () => {
         () => defaultDriverOptions
     );
 
-    it.skipIf(isMacOSInGitHubActions)("should have the atlas-local-list-deployments tool", async () => {
+    it.skipIf(isMacOSInGitHubActions)("should have the atlas-local-list-deployments tool", async ({ signal }) => {
+        await waitUntilMcpClientIsSet(integration.mcpServer(), signal);
+
         const { tools } = await integration.mcpClient().listTools();
         const listDeployments = tools.find((tool) => tool.name === "atlas-local-list-deployments");
         expectDefined(listDeployments);
     });
 
     it.skipIf(!isMacOSInGitHubActions)(
         "[MacOS in GitHub Actions] should not have the atlas-local-list-deployments tool",
-        async () => {
+        async ({ signal }) => {
+            await waitUntilMcpClientIsSet(integration.mcpServer(), signal);
             const { tools } = await integration.mcpClient().listTools();
             const listDeployments = tools.find((tool) => tool.name === "atlas-local-list-deployments");
             expect(listDeployments).toBeUndefined();
         }
     );
 
-    it.skipIf(isMacOSInGitHubActions)("should have correct metadata", async () => {
+    it.skipIf(isMacOSInGitHubActions)("should have correct metadata", async ({ signal }) => {
+        await waitUntilMcpClientIsSet(integration.mcpServer(), signal);
         const { tools } = await integration.mcpClient().listTools();
         const listDeployments = tools.find((tool) => tool.name === "atlas-local-list-deployments");
         expectDefined(listDeployments);
@@ -41,7 +46,9 @@ describe("atlas-local-list-deployments", () => {
         expect(listDeployments.inputSchema.properties).toEqual({});
     });
 
-    it.skipIf(isMacOSInGitHubActions)("should not crash when calling the tool", async () => {
+    it.skipIf(isMacOSInGitHubActions)("should not crash when calling the tool", async ({ signal }) => {
+        await waitUntilMcpClientIsSet(integration.mcpServer(), signal);
+
         const response = await integration.mcpClient().callTool({
             name: "atlas-local-list-deployments",
             arguments: {},