Use hooks, clean up test setup a little

nirinchev · nirinchev · commit 6f330a363792 · 2025-04-14T13:43:02.000+02:00
diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts
@@ -7,65 +7,100 @@ import fs from "fs/promises";
 import { Session } from "../../src/session.js";
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 
-export async function setupIntegrationTest(): Promise<{
-    client: Client;
-    server: Server;
-    teardown: () => Promise<void>;
-}> {
-    const clientTransport = new InMemoryTransport();
-    const serverTransport = new InMemoryTransport();
-
-    await serverTransport.start();
-    await clientTransport.start();
-
-    clientTransport.output.pipeTo(serverTransport.input);
-    serverTransport.output.pipeTo(clientTransport.input);
-
-    const client = new Client(
-        {
-            name: "test-client",
-            version: "1.2.3",
-        },
-        {
-            capabilities: {},
-        }
-    );
-
-    const server = new Server({
-        mcpServer: new McpServer({
-            name: "test-server",
-            version: "1.2.3",
-        }),
-        session: new Session(),
+export function jestTestMCPClient(): () => Client {
+    let client: Client | undefined;
+    let server: Server | undefined;
+
+    beforeEach(async () => {
+        const clientTransport = new InMemoryTransport();
+        const serverTransport = new InMemoryTransport();
+
+        await serverTransport.start();
+        await clientTransport.start();
+
+        clientTransport.output.pipeTo(serverTransport.input);
+        serverTransport.output.pipeTo(clientTransport.input);
+
+        client = new Client(
+            {
+                name: "test-client",
+                version: "1.2.3",
+            },
+            {
+                capabilities: {},
+            }
+        );
+
+        server = new Server({
+            mcpServer: new McpServer({
+                name: "test-server",
+                version: "1.2.3",
+            }),
+            session: new Session(),
+        });
+        await server.connect(serverTransport);
+        await client.connect(clientTransport);
+    });
+
+    afterEach(async () => {
+        await client?.close();
+        client = undefined;
+
+        await server?.close();
+        server = undefined;
     });
-    await server.connect(serverTransport);
-    await client.connect(clientTransport);
-
-    return {
-        client,
-        server,
-        teardown: async () => {
-            await client.close();
-            await server.close();
-        },
+
+    return () => {
+        if (!client) {
+            throw new Error("beforeEach() hook not ran yet");
+        }
+
+        return client;
     };
 }
 
-export async function runMongoDB(): Promise<runner.MongoCluster> {
-    const tmpDir = path.join(__dirname, "..", "tmp");
-    await fs.mkdir(tmpDir, { recursive: true });
+export function jestTestCluster(): () => runner.MongoCluster {
+    let cluster: runner.MongoCluster | undefined;
 
-    try {
-        const cluster = await MongoCluster.start({
-            tmpDir: path.join(tmpDir, "mongodb-runner", "dbs"),
-            logDir: path.join(tmpDir, "mongodb-runner", "logs"),
-            topology: "standalone",
-        });
+    function runMongodb() {}
+
+    beforeAll(async function () {
+        // Downloading Windows executables in CI takes a long time because
+        // they include debug symbols...
+        const tmpDir = path.join(__dirname, "..", "tmp");
+        await fs.mkdir(tmpDir, { recursive: true });
+
+        // On Windows, we may have a situation where mongod.exe is not fully released by the OS
+        // before we attempt to run it again, so we add a retry.
+        const dbsDir = path.join(tmpDir, "mongodb-runner", `dbs`);
+        for (let i = 0; i < 10; i++) {
+            try {
+                cluster = await MongoCluster.start({
+                    tmpDir: dbsDir,
+                    logDir: path.join(tmpDir, "mongodb-runner", "logs"),
+                    topology: "standalone",
+                });
+
+                return;
+            } catch (err) {
+                console.error(`Failed to start cluster in ${dbsDir}, attempt ${i}: ${err}`);
+                await new Promise((resolve) => setTimeout(resolve, 1000));
+            }
+        }
+    }, 120_000);
+
+    afterAll(async function () {
+        await cluster?.close();
+        cluster = undefined;
+    });
+
+    return () => {
+        if (!cluster) {
+            throw new Error("beforeAll() hook not ran yet");
+        }
 
         return cluster;
-    } catch (err) {
-        throw err;
-    }
+    };
 }
 
 export function getResponseContent(content: unknown): string {
diff --git a/tests/integration/server.test.ts b/tests/integration/server.test.ts
@@ -1,39 +1,29 @@
-import { Client } from "@modelcontextprotocol/sdk/client/index.js";
-import { setupIntegrationTest } from "./helpers.js";
+import { jestTestMCPClient } from "./helpers.js";
 
 describe("Server integration test", () => {
-    let client: Client;
-    let teardown: () => Promise<void>;
-
-    beforeEach(async () => {
-        ({ client, teardown } = await setupIntegrationTest());
-    });
-
-    afterEach(async () => {
-        await teardown();
-    });
+    const client = jestTestMCPClient();
 
     describe("list capabilities", () => {
         it("should return positive number of tools", async () => {
-            const tools = await client.listTools();
+            const tools = await client().listTools();
             expect(tools).toBeDefined();
             expect(tools.tools.length).toBeGreaterThan(0);
         });
 
         it("should return no resources", async () => {
-            await expect(() => client.listResources()).rejects.toMatchObject({
+            await expect(() => client().listResources()).rejects.toMatchObject({
                 message: "MCP error -32601: Method not found",
             });
         });
 
         it("should return no prompts", async () => {
-            await expect(() => client.listPrompts()).rejects.toMatchObject({
+            await expect(() => client().listPrompts()).rejects.toMatchObject({
                 message: "MCP error -32601: Method not found",
             });
         });
 
         it("should return capabilities", async () => {
-            const capabilities = client.getServerCapabilities();
+            const capabilities = client().getServerCapabilities();
             expect(capabilities).toBeDefined();
             expect(capabilities?.completions).toBeUndefined();
             expect(capabilities?.experimental).toBeUndefined();
diff --git a/tests/integration/tools/mongodb/metadata/connect.test.ts b/tests/integration/tools/mongodb/metadata/connect.test.ts
@@ -1,33 +1,13 @@
-import { Client } from "@modelcontextprotocol/sdk/client/index.js";
-import { runMongoDB, setupIntegrationTest, getResponseContent } from "../../../helpers.js";
-import runner from "mongodb-runner";
+import { getResponseContent, jestTestMCPClient, jestTestCluster } from "../../../helpers.js";
 
 import config from "../../../../../src/config.js";
 
 describe("Connect tool", () => {
-    let client: Client;
-    let serverClientTeardown: () => Promise<void>;
-
-    let cluster: runner.MongoCluster;
-
-    beforeAll(async () => {
-        cluster = await runMongoDB();
-    }, 60_000);
-
-    beforeEach(async () => {
-        ({ client, teardown: serverClientTeardown } = await setupIntegrationTest());
-    });
-
-    afterEach(async () => {
-        await serverClientTeardown?.();
-    });
-
-    afterAll(async () => {
-        await cluster.close();
-    });
+    const client = jestTestMCPClient();
+    const cluster = jestTestCluster();
 
     it("should have correct metadata", async () => {
-        const { tools } = await client.listTools();
+        const { tools } = await client().listTools();
         const connectTool = tools.find((tool) => tool.name === "connect")!;
         expect(connectTool).toBeDefined();
         expect(connectTool.description).toBe("Connect to a MongoDB instance");
@@ -50,7 +30,7 @@ describe("Connect tool", () => {
     describe("with default config", () => {
         describe("without connection string", () => {
             it("prompts for connection string", async () => {
-                const response = await client.callTool({ name: "connect", arguments: {} });
+                const response = await client().callTool({ name: "connect", arguments: {} });
                 const content = getResponseContent(response.content);
                 expect(content).toContain("No connection details provided");
                 expect(content).toContain("mongodb://localhost:27017");
@@ -59,19 +39,19 @@ describe("Connect tool", () => {
 
         describe("with connection string", () => {
             it("connects to the database", async () => {
-                const response = await client.callTool({
+                const response = await client().callTool({
                     name: "connect",
-                    arguments: { connectionStringOrClusterName: cluster.connectionString },
+                    arguments: { connectionStringOrClusterName: cluster().connectionString },
                 });
                 const content = getResponseContent(response.content);
                 expect(content).toContain("Successfully connected");
-                expect(content).toContain(cluster.connectionString);
+                expect(content).toContain(cluster().connectionString);
             });
         });
 
         describe("with invalid connection string", () => {
             it("returns error message", async () => {
-                const response = await client.callTool({
+                const response = await client().callTool({
                     name: "connect",
                     arguments: { connectionStringOrClusterName: "mongodb://localhost:12345" },
                 });
@@ -83,19 +63,19 @@ describe("Connect tool", () => {
 
     describe("with connection string in config", () => {
         beforeEach(async () => {
-            config.connectionString = cluster.connectionString;
+            config.connectionString = cluster().connectionString;
         });
 
         it("uses the connection string from config", async () => {
-            const response = await client.callTool({ name: "connect", arguments: {} });
+            const response = await client().callTool({ name: "connect", arguments: {} });
             const content = getResponseContent(response.content);
             expect(content).toContain("Successfully connected");
-            expect(content).toContain(cluster.connectionString);
+            expect(content).toContain(cluster().connectionString);
         });
 
         it("prefers connection string from arguments", async () => {
-            const newConnectionString = `${cluster.connectionString}?appName=foo-bar`;
-            const response = await client.callTool({
+            const newConnectionString = `${cluster().connectionString}?appName=foo-bar`;
+            const response = await client().callTool({
                 name: "connect",
                 arguments: { connectionStringOrClusterName: newConnectionString },
             });
diff --git a/tests/integration/tools/mongodb/metadata/listDatabases.test.ts b/tests/integration/tools/mongodb/metadata/listDatabases.test.ts
@@ -1,33 +1,13 @@
-import { Client } from "@modelcontextprotocol/sdk/client/index.js";
-import { runMongoDB, setupIntegrationTest, getResponseElements, connect } from "../../../helpers.js";
-import runner from "mongodb-runner";
+import { getResponseElements, connect, jestTestCluster, jestTestMCPClient } from "../../../helpers.js";
 import { MongoClient } from "mongodb";
 import { toIncludeSameMembers } from "jest-extended";
 
 describe("listDatabases tool", () => {
-    let client: Client;
-    let serverClientTeardown: () => Promise<void>;
-
-    let cluster: runner.MongoCluster;
-
-    beforeAll(async () => {
-        cluster = await runMongoDB();
-    }, 60_000);
-
-    beforeEach(async () => {
-        ({ client, teardown: serverClientTeardown } = await setupIntegrationTest());
-    });
-
-    afterEach(async () => {
-        await serverClientTeardown?.();
-    });
-
-    afterAll(async () => {
-        await cluster.close();
-    });
+    const client = jestTestMCPClient();
+    const cluster = jestTestCluster();
 
     it("should have correct metadata", async () => {
-        const { tools } = await client.listTools();
+        const { tools } = await client().listTools();
         const listDatabases = tools.find((tool) => tool.name === "list-databases")!;
         expect(listDatabases).toBeDefined();
         expect(listDatabases.description).toBe("List all databases for a MongoDB connection");
@@ -40,8 +20,8 @@ describe("listDatabases tool", () => {
 
     describe("with no preexisting databases", () => {
         it("returns only the system databases", async () => {
-            await connect(client, cluster);
-            const response = await client.callTool({ name: "list-databases", arguments: {} });
+            await connect(client(), cluster());
+            const response = await client().callTool({ name: "list-databases", arguments: {} });
             const dbNames = getDbNames(response.content);
 
             expect(dbNames).toIncludeSameMembers(["admin", "config", "local"]);
@@ -50,14 +30,14 @@ describe("listDatabases tool", () => {
 
     describe("with preexisting databases", () => {
         it("returns their names and sizes", async () => {
-            const mongoClient = new MongoClient(cluster.connectionString);
+            const mongoClient = new MongoClient(cluster().connectionString);
             await mongoClient.db("foo").collection("bar").insertOne({ test: "test" });
             await mongoClient.db("baz").collection("qux").insertOne({ test: "test" });
             await mongoClient.close();
 
-            await connect(client, cluster);
+            await connect(client(), cluster());
 
-            const response = await client.callTool({ name: "list-databases", arguments: {} });
+            const response = await client().callTool({ name: "list-databases", arguments: {} });
             const dbNames = getDbNames(response.content);
             expect(dbNames).toIncludeSameMembers(["admin", "config", "local", "foo", "baz"]);
         });
