chore: refactor integration test setup

tests/integration/helpers.ts

@@ -18,9 +18,17 @@ interface ParameterInfo {
 
 type ToolInfo = Awaited<ReturnType<Client["listTools"]>>["tools"][number];
 
-export function jestTestMCPClient(): () => Client {
-    let client: Client | undefined;
-    let server: Server | undefined;
+export function setupIntegrationTest(): {
+    mcpClient: () => Client;
+    mongoClient: () => MongoClient;
+    connectionString: () => string;
+    connectMcpClient: () => Promise<void>;
+} {
+    let mongoCluster: runner.MongoCluster | undefined;
+    let mongoClient: MongoClient | undefined;
+
+    let mcpClient: Client | undefined;
+    let mcpServer: Server | undefined;
 
     beforeEach(async () => {
         const clientTransport = new InMemoryTransport();
@@ -32,7 +40,7 @@ export function jestTestMCPClient(): () => Client {
         clientTransport.output.pipeTo(serverTransport.input);
         serverTransport.output.pipeTo(clientTransport.input);
 
-        client = new Client(
+        mcpClient = new Client(
             {
                 name: "test-client",
                 version: "1.2.3",
@@ -42,41 +50,26 @@ export function jestTestMCPClient(): () => Client {
             }
         );
 
-        server = new Server({
+        mcpServer = new Server({
             mcpServer: new McpServer({
                 name: "test-server",
                 version: "1.2.3",
             }),
             session: new Session(),
         });
-        await server.connect(serverTransport);
-        await client.connect(clientTransport);
+        await mcpServer.connect(serverTransport);
+        await mcpClient.connect(clientTransport);
     });
 
     afterEach(async () => {
-        await client?.close();
-        client = undefined;
+        await mcpClient?.close();
+        mcpClient = undefined;
 
-        await server?.close();
-        server = undefined;
-    });
+        await mcpServer?.close();
+        mcpServer = undefined;
 
-    return () => {
-        if (!client) {
-            throw new Error("beforeEach() hook not ran yet");
-        }
-
-        return client;
-    };
-}
-
-export function jestTestCluster(): () => { connectionString: string; getClient: () => MongoClient } {
-    let cluster: runner.MongoCluster | undefined;
-    let client: MongoClient | undefined;
-
-    afterEach(async () => {
-        await client?.close();
-        client = undefined;
+        await mongoClient?.close();
+        mongoClient = undefined;
     });
 
     beforeAll(async function () {
@@ -90,7 +83,7 @@ export function jestTestCluster(): () => { connectionString: string; getClient:
         let dbsDir = path.join(tmpDir, "mongodb-runner", "dbs");
         for (let i = 0; i < 10; i++) {
             try {
-                cluster = await MongoCluster.start({
+                mongoCluster = await MongoCluster.start({
                     tmpDir: dbsDir,
                     logDir: path.join(tmpDir, "mongodb-runner", "logs"),
                     topology: "standalone",
@@ -116,25 +109,41 @@ export function jestTestCluster(): () => { connectionString: string; getClient:
     }, 120_000);
 
     afterAll(async function () {
-        await cluster?.close();
-        cluster = undefined;
+        await mongoCluster?.close();
+        mongoCluster = undefined;
     });
 
-    return () => {
-        if (!cluster) {
+    const getMcpClient = () => {
+        if (!mcpClient) {
+            throw new Error("beforeEach() hook not ran yet");
+        }
+
+        return mcpClient;
+    };
+
+    const getConnectionString = () => {
+        if (!mongoCluster) {
             throw new Error("beforeAll() hook not ran yet");
         }
 
-        return {
-            connectionString: cluster.connectionString,
-            getClient: () => {
-                if (!client) {
-                    client = new MongoClient(cluster!.connectionString);
-                }
+        return mongoCluster.connectionString;
+    };
 
-                return client;
-            },
-        };
+    return {
+        mcpClient: getMcpClient,
+        mongoClient: () => {
+            if (!mongoClient) {
+                mongoClient = new MongoClient(getConnectionString());
+            }
+            return mongoClient;
+        },
+        connectionString: getConnectionString,
+        connectMcpClient: async () => {
+            await getMcpClient().callTool({
+                name: "connect",
+                arguments: { connectionStringOrClusterName: getConnectionString() },
+            });
+        },
     };
 }
 
@@ -157,10 +166,10 @@ export function getResponseElements(content: unknown): { type: string; text: str
     return response;
 }
 
-export async function connect(client: Client, cluster: runner.MongoCluster): Promise<void> {
+export async function connect(client: Client, connectionString: string): Promise<void> {
     await client.callTool({
         name: "connect",
-        arguments: { connectionStringOrClusterName: cluster.connectionString },
+        arguments: { connectionStringOrClusterName: connectionString },
     });
 }
 

tests/integration/server.test.ts

@@ -1,29 +1,29 @@
-import { jestTestMCPClient } from "./helpers.js";
+import { setupIntegrationTest } from "./helpers";
 
 describe("Server integration test", () => {
-    const client = jestTestMCPClient();
+    const integration = setupIntegrationTest();
 
     describe("list capabilities", () => {
         it("should return positive number of tools", async () => {
-            const tools = await client().listTools();
+            const tools = await integration.mcpClient().listTools();
             expect(tools).toBeDefined();
             expect(tools.tools.length).toBeGreaterThan(0);
         });
 
         it("should return no resources", async () => {
-            await expect(() => client().listResources()).rejects.toMatchObject({
+            await expect(() => integration.mcpClient().listResources()).rejects.toMatchObject({
                 message: "MCP error -32601: Method not found",
             });
         });
 
         it("should return no prompts", async () => {
-            await expect(() => client().listPrompts()).rejects.toMatchObject({
+            await expect(() => integration.mcpClient().listPrompts()).rejects.toMatchObject({
                 message: "MCP error -32601: Method not found",
             });
         });
 
         it("should return capabilities", async () => {
-            const capabilities = client().getServerCapabilities();
+            const capabilities = integration.mcpClient().getServerCapabilities();
             expect(capabilities).toBeDefined();
             expect(capabilities?.completions).toBeUndefined();
             expect(capabilities?.experimental).toBeUndefined();

tests/integration/tools/mongodb/create/createCollection.test.ts

@@ -1,21 +1,18 @@
 import {
-    connect,
-    jestTestCluster,
-    jestTestMCPClient,
     getResponseContent,
     validateParameters,
     dbOperationParameters,
+    setupIntegrationTest,
 } from "../../../helpers.js";
 import { toIncludeSameMembers } from "jest-extended";
 import { McpError } from "@modelcontextprotocol/sdk/types.js";
 import { ObjectId } from "bson";
 
 describe("createCollection tool", () => {
-    const client = jestTestMCPClient();
-    const cluster = jestTestCluster();
+    const integration = setupIntegrationTest();
 
     it("should have correct metadata", async () => {
-        const { tools } = await client().listTools();
+        const { tools } = await integration.mcpClient().listTools();
         const listCollections = tools.find((tool) => tool.name === "create-collection")!;
         expect(listCollections).toBeDefined();
         expect(listCollections.description).toBe(
@@ -34,9 +31,9 @@ describe("createCollection tool", () => {
         ];
         for (const arg of args) {
             it(`throws a schema error for: ${JSON.stringify(arg)}`, async () => {
-                await connect(client(), cluster());
+                await integration.connectMcpClient();
                 try {
-                    await client().callTool({ name: "create-collection", arguments: arg });
+                    await integration.mcpClient().callTool({ name: "create-collection", arguments: arg });
                     expect.fail("Expected an error to be thrown");
                 } catch (error) {
                     expect(error).toBeInstanceOf(McpError);
@@ -50,12 +47,12 @@ describe("createCollection tool", () => {
 
     describe("with non-existent database", () => {
         it("creates a new collection", async () => {
-            const mongoClient = cluster().getClient();
+            const mongoClient = integration.mongoClient();
             let collections = await mongoClient.db("foo").listCollections().toArray();
             expect(collections).toHaveLength(0);
 
-            await connect(client(), cluster());
-            const response = await client().callTool({
+            await integration.connectMcpClient();
+            const response = await integration.mcpClient().callTool({
                 name: "create-collection",
                 arguments: { database: "foo", collection: "bar" },
             });
@@ -75,13 +72,13 @@ describe("createCollection tool", () => {
         });
 
         it("creates new collection", async () => {
-            const mongoClient = cluster().getClient();
+            const mongoClient = integration.mongoClient();
             await mongoClient.db(dbName).createCollection("collection1");
             let collections = await mongoClient.db(dbName).listCollections().toArray();
             expect(collections).toHaveLength(1);
 
-            await connect(client(), cluster());
-            const response = await client().callTool({
+            await integration.connectMcpClient();
+            const response = await integration.mcpClient().callTool({
                 name: "create-collection",
                 arguments: { database: dbName, collection: "collection2" },
             });
@@ -93,15 +90,15 @@ describe("createCollection tool", () => {
         });
 
         it("does nothing if collection already exists", async () => {
-            const mongoClient = cluster().getClient();
+            const mongoClient = integration.mongoClient();
             await mongoClient.db(dbName).collection("collection1").insertOne({});
             let collections = await mongoClient.db(dbName).listCollections().toArray();
             expect(collections).toHaveLength(1);
             let documents = await mongoClient.db(dbName).collection("collection1").find({}).toArray();
             expect(documents).toHaveLength(1);
 
-            await connect(client(), cluster());
-            const response = await client().callTool({
+            await integration.connectMcpClient();
+            const response = await integration.mcpClient().callTool({
                 name: "create-collection",
                 arguments: { database: dbName, collection: "collection1" },
             });

tests/integration/tools/mongodb/metadata/connect.test.ts

@@ -1,13 +1,12 @@
-import { getResponseContent, jestTestMCPClient, jestTestCluster, validateParameters } from "../../../helpers.js";
+import { getResponseContent, validateParameters, setupIntegrationTest } from "../../../helpers.js";
 
 import config from "../../../../../src/config.js";
 
 describe("Connect tool", () => {
-    const client = jestTestMCPClient();
-    const cluster = jestTestCluster();
+    const integration = setupIntegrationTest();
 
     it("should have correct metadata", async () => {
-        const { tools } = await client().listTools();
+        const { tools } = await integration.mcpClient().listTools();
         const connectTool = tools.find((tool) => tool.name === "connect")!;
         expect(connectTool).toBeDefined();
         expect(connectTool.description).toBe("Connect to a MongoDB instance");
@@ -25,7 +24,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 integration.mcpClient().callTool({ name: "connect", arguments: {} });
                 const content = getResponseContent(response.content);
                 expect(content).toContain("No connection details provided");
                 expect(content).toContain("mongodb://localhost:27017");
@@ -34,19 +33,19 @@ describe("Connect tool", () => {
 
         describe("with connection string", () => {
             it("connects to the database", async () => {
-                const response = await client().callTool({
+                const response = await integration.mcpClient().callTool({
                     name: "connect",
-                    arguments: { connectionStringOrClusterName: cluster().connectionString },
+                    arguments: { connectionStringOrClusterName: integration.connectionString() },
                 });
                 const content = getResponseContent(response.content);
                 expect(content).toContain("Successfully connected");
-                expect(content).toContain(cluster().connectionString);
+                expect(content).toContain(integration.connectionString());
             });
         });
 
         describe("with invalid connection string", () => {
             it("returns error message", async () => {
-                const response = await client().callTool({
+                const response = await integration.mcpClient().callTool({
                     name: "connect",
                     arguments: { connectionStringOrClusterName: "mongodb://localhost:12345" },
                 });
@@ -61,19 +60,19 @@ describe("Connect tool", () => {
 
     describe("with connection string in config", () => {
         beforeEach(async () => {
-            config.connectionString = cluster().connectionString;
+            config.connectionString = integration.connectionString();
         });
 
         it("uses the connection string from config", async () => {
-            const response = await client().callTool({ name: "connect", arguments: {} });
+            const response = await integration.mcpClient().callTool({ name: "connect", arguments: {} });
             const content = getResponseContent(response.content);
             expect(content).toContain("Successfully connected");
-            expect(content).toContain(cluster().connectionString);
+            expect(content).toContain(integration.connectionString());
         });
 
         it("prefers connection string from arguments", async () => {
-            const newConnectionString = `${cluster().connectionString}?appName=foo-bar`;
-            const response = await client().callTool({
+            const newConnectionString = `${integration.connectionString()}?appName=foo-bar`;
+            const response = await integration.mcpClient().callTool({
                 name: "connect",
                 arguments: { connectionStringOrClusterName: newConnectionString },
             });
@@ -84,7 +83,7 @@ describe("Connect tool", () => {
 
         describe("when the arugment connection string is invalid", () => {
             it("suggests the config connection string if set", async () => {
-                const response = await client().callTool({
+                const response = await integration.mcpClient().callTool({
                     name: "connect",
                     arguments: { connectionStringOrClusterName: "mongodb://localhost:12345" },
                 });
@@ -97,7 +96,7 @@ describe("Connect tool", () => {
 
             it("returns error message if the config connection string matches the argument", async () => {
                 config.connectionString = "mongodb://localhost:12345";
-                const response = await client().callTool({
+                const response = await integration.mcpClient().callTool({
                     name: "connect",
                     arguments: { connectionStringOrClusterName: "mongodb://localhost:12345" },
                 });