fix integration tests

tests/integration/helpers.ts

@@ -19,10 +19,10 @@ 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/atlas-local";
 import { Elicitation } from "../../src/elicitation.js";
 import type { MockClientCapabilities, createMockElicitInput } from "../utils/elicitationMocks.js";
 import { VectorSearchEmbeddingsManager } from "../../src/common/search/vectorSearchEmbeddingsManager.js";
+import { defaultCreateAtlasLocalClient } from "../../src/common/atlasLocal.js";
 
 export const driverOptions = setupDriverConfig({
  config,
@@ -115,6 +115,7 @@ export function setupIntegrationTest(
  connectionManager,
  keychain: new Keychain(),
  vectorSearchEmbeddingsManager: new VectorSearchEmbeddingsManager(userConfig, connectionManager),
+ atlasLocalClient: await defaultCreateAtlasLocalClient(),
  });
 
  // Mock hasValidAccessToken for tests
@@ -408,37 +409,6 @@ export function waitUntil<T extends ConnectionState>(
  });
 }
 
-export function waitUntilAtlasLocalClientIsSet(
- 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/tools/atlas-local/connectDeployment.test.ts

@@ -6,7 +6,6 @@ import {
  getResponseElements,
  setupIntegrationTest,
  validateToolMetadata,
- waitUntilAtlasLocalClientIsSet,
 } from "../../helpers.js";
 import { afterEach, describe, expect, it } from "vitest";
 
@@ -19,10 +18,6 @@ const integration = setupIntegrationTest(
 // Docker is not available on macOS in GitHub Actions
 // That's why we skip the tests on macOS in GitHub Actions
 describe.skipIf(isMacOSInGitHubActions)("atlas-local-connect-deployment", () => {
- beforeEach(async ({ signal }) => {
- await waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal);
- });
-
  validateToolMetadata(integration, "atlas-local-connect-deployment", "Connect to a MongoDB Atlas Local deployment", [
  {
  name: "deploymentName",
@@ -56,9 +51,7 @@ describe.skipIf(isMacOSInGitHubActions)("atlas-local-connect-deployment with dep
  let deploymentName: string = "";
  let deploymentNamesToCleanup: string[] = [];
 
- beforeEach(async ({ signal }) => {
- await waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal);
-
+ beforeEach(async () => {
  // Create deployments
  deploymentName = `test-deployment-1-${Date.now()}`;
  deploymentNamesToCleanup.push(deploymentName);
@@ -103,10 +96,8 @@ describe.skipIf(isMacOSInGitHubActions)("atlas-local-connect-deployment with dep
 });
 
 describe.skipIf(!isMacOSInGitHubActions)("atlas-local-connect-deployment [MacOS in GitHub Actions]", () => {
- it("should not have the atlas-local-connect-deployment tool", async ({ signal }) => {
+ it("should not have the atlas-local-connect-deployment tool", async () => {
  // This should throw an error because the client is not set within the timeout of 5 seconds (default)
- await expect(waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal)).rejects.toThrow();
-
  const { tools } = await integration.mcpClient().listTools();
  const connectDeployment = tools.find((tool) => tool.name === "atlas-local-connect-deployment");
  expect(connectDeployment).toBeUndefined();

tests/integration/tools/atlas-local/createDeployment.test.ts

@@ -4,7 +4,6 @@ import {
  expectDefined,
  getResponseElements,
  setupIntegrationTest,
- waitUntilAtlasLocalClientIsSet,
 } from "../../helpers.js";
 import { afterEach, describe, expect, it } from "vitest";
 
@@ -34,28 +33,22 @@ describe("atlas-local-create-deployment", () => {
  () => defaultDriverOptions
  );
 
- it.skipIf(isMacOSInGitHubActions)("should have the atlas-local-create-deployment tool", async ({ signal }) => {
- await waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal);
-
+ it.skipIf(isMacOSInGitHubActions)("should have the atlas-local-create-deployment tool", async () => {
  const { tools } = await integration.mcpClient().listTools();
  const createDeployment = tools.find((tool) => tool.name === "atlas-local-create-deployment");
  expectDefined(createDeployment);
  });
 
  it.skipIf(!isMacOSInGitHubActions)(
  "[MacOS in GitHub Actions] should not have the atlas-local-create-deployment tool",
- async ({ signal }) => {
- // This should throw an error because the client is not set within the timeout of 5 seconds (default)
- await expect(waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal)).rejects.toThrow();
-
+ async () => {
  const { tools } = await integration.mcpClient().listTools();
  const createDeployment = tools.find((tool) => tool.name === "atlas-local-create-deployment");
  expect(createDeployment).toBeUndefined();
  }
  );
 
- it.skipIf(isMacOSInGitHubActions)("should have correct metadata", async ({ signal }) => {
- await waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal);
+ it.skipIf(isMacOSInGitHubActions)("should have correct metadata", async () => {
  const { tools } = await integration.mcpClient().listTools();
  const createDeployment = tools.find((tool) => tool.name === "atlas-local-create-deployment");
  expectDefined(createDeployment);
@@ -64,8 +57,7 @@ describe("atlas-local-create-deployment", () => {
  expect(createDeployment.inputSchema.properties).toHaveProperty("deploymentName");
  });
 
- it.skipIf(isMacOSInGitHubActions)("should create a deployment when calling the tool", async ({ signal }) => {
- await waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal);
+ it.skipIf(isMacOSInGitHubActions)("should create a deployment when calling the tool", async () => {
  const deploymentName = `test-deployment-${Date.now()}`;
 
  // Check that deployment doesn't exist before creation
@@ -97,9 +89,7 @@ describe("atlas-local-create-deployment", () => {
 
  it.skipIf(isMacOSInGitHubActions)(
  "should return an error when creating a deployment that already exists",
- async ({ signal }) => {
- await waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal);
-
+ async () => {
  // Create a deployment
  const deploymentName = `test-deployment-${Date.now()}`;
  deploymentNamesToCleanup.push(deploymentName);
@@ -119,9 +109,7 @@ describe("atlas-local-create-deployment", () => {
  }
  );
 
- it.skipIf(isMacOSInGitHubActions)("should create a deployment with the correct name", async ({ signal }) => {
- await waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal);
-
+ it.skipIf(isMacOSInGitHubActions)("should create a deployment with the correct name", async () => {
  // Create a deployment
  const deploymentName = `test-deployment-${Date.now()}`;
  deploymentNamesToCleanup.push(deploymentName);
@@ -147,9 +135,7 @@ describe("atlas-local-create-deployment", () => {
  expect(elements[1]?.text ?? "").toContain("Running");
  });
 
- it.skipIf(isMacOSInGitHubActions)("should create a deployment when name is not provided", async ({ signal }) => {
- await waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal);
-
+ it.skipIf(isMacOSInGitHubActions)("should create a deployment when name is not provided", async () => {
  // Create a deployment
  const createResponse = await integration.mcpClient().callTool({
  name: "atlas-local-create-deployment",

tests/integration/tools/atlas-local/deleteDeployment.test.ts

@@ -4,7 +4,6 @@ import {
  expectDefined,
  getResponseElements,
  setupIntegrationTest,
- waitUntilAtlasLocalClientIsSet,
 } from "../../helpers.js";
 import { describe, expect, it } from "vitest";
 
@@ -18,28 +17,22 @@ describe("atlas-local-delete-deployment", () => {
  () => defaultDriverOptions
  );
 
- it.skipIf(isMacOSInGitHubActions)("should have the atlas-local-delete-deployment tool", async ({ signal }) => {
- await waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal);
-
+ it.skipIf(isMacOSInGitHubActions)("should have the atlas-local-delete-deployment tool", async () => {
  const { tools } = await integration.mcpClient().listTools();
  const deleteDeployment = tools.find((tool) => tool.name === "atlas-local-delete-deployment");
  expectDefined(deleteDeployment);
  });
 
  it.skipIf(!isMacOSInGitHubActions)(
  "[MacOS in GitHub Actions] should not have the atlas-local-delete-deployment tool",
- async ({ signal }) => {
- // This should throw an error because the client is not set within the timeout of 5 seconds (default)
- await expect(waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal)).rejects.toThrow();
-
+ async () => {
  const { tools } = await integration.mcpClient().listTools();
  const deleteDeployment = tools.find((tool) => tool.name === "atlas-local-delete-deployment");
  expect(deleteDeployment).toBeUndefined();
  }
  );
 
- it.skipIf(isMacOSInGitHubActions)("should have correct metadata", async ({ signal }) => {
- await waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal);
+ it.skipIf(isMacOSInGitHubActions)("should have correct metadata", async () => {
  const { tools } = await integration.mcpClient().listTools();
  const deleteDeployment = tools.find((tool) => tool.name === "atlas-local-delete-deployment");
  expectDefined(deleteDeployment);
@@ -50,8 +43,7 @@ describe("atlas-local-delete-deployment", () => {
 
  it.skipIf(isMacOSInGitHubActions)(
  "should return 'no such container' error when deployment to delete does not exist",
- async ({ signal }) => {
- await waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal);
+ async () => {
  const deploymentName = "non-existent";
 
  const response = await integration.mcpClient().callTool({
@@ -66,8 +58,7 @@ describe("atlas-local-delete-deployment", () => {
  }
  );
 
- it.skipIf(isMacOSInGitHubActions)("should delete a deployment when calling the tool", async ({ signal }) => {
- await waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal);
+ it.skipIf(isMacOSInGitHubActions)("should delete a deployment when calling the tool", async () => {
  // Create a deployment
  const deploymentName = `test-deployment-${Date.now()}`;
  await integration.mcpClient().callTool({

tests/integration/tools/atlas-local/listDeployments.test.ts

@@ -4,7 +4,6 @@ import {
  expectDefined,
  getResponseElements,
  setupIntegrationTest,
- waitUntilAtlasLocalClientIsSet,
 } from "../../helpers.js";
 import { describe, expect, it } from "vitest";
 
@@ -18,28 +17,22 @@ describe("atlas-local-list-deployments", () => {
  () => defaultDriverOptions
  );
 
- it.skipIf(isMacOSInGitHubActions)("should have the atlas-local-list-deployments tool", async ({ signal }) => {
- await waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal);
-
+ it.skipIf(isMacOSInGitHubActions)("should have the atlas-local-list-deployments tool", async () => {
  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 ({ signal }) => {
- // This should throw an error because the client is not set within the timeout of 5 seconds (default)
- await expect(waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal)).rejects.toThrow();
-
+ async () => {
  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 ({ signal }) => {
- await waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal);
+ it.skipIf(isMacOSInGitHubActions)("should have correct metadata", async () => {
  const { tools } = await integration.mcpClient().listTools();
  const listDeployments = tools.find((tool) => tool.name === "atlas-local-list-deployments");
  expectDefined(listDeployments);
@@ -48,9 +41,7 @@ describe("atlas-local-list-deployments", () => {
  expect(listDeployments.inputSchema.properties).toEqual({});
  });
 
- it.skipIf(isMacOSInGitHubActions)("should not crash when calling the tool", async ({ signal }) => {
- await waitUntilAtlasLocalClientIsSet(integration.mcpServer(), signal);
-
+ it.skipIf(isMacOSInGitHubActions)("should not crash when calling the tool", async () => {
  const response = await integration.mcpClient().callTool({
  name: "atlas-local-list-deployments",
  arguments: {},