chore: replace mock mcp client with real (mockable) mcp client

When writing test cases, I realized that it is too much duplicated effort to write and maintain mocks. So instead of having only a mocked mcp client, this commit introduces a real mcp client that talks to our mcp server and is still mockable.

We are now setting up real MCP client with test data in mongodb database spun up for test suites. Mocking is still an option but we likely never feel the need for that.

Author: himanshusinghs

tests/accuracy/sdk/accuracy-testing-client.ts

@@ -0,0 +1,76 @@
+import path from "path";
+import { v4 as uuid } from "uuid";
+import { fileURLToPath } from "url";
+import { experimental_createMCPClient as createMCPClient, tool as createVercelTool } from "ai";
+import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
+
+import { ToolCall } from "./accuracy-scorers.js";
+
+const __dirname = fileURLToPath(import.meta.url);
+const distPath = path.join(__dirname, "..", "..", "..", "..", "dist");
+const cliScriptPath = path.join(distPath, "index.js");
+
+type ToolResultGeneratorFn = (...parameters: unknown[]) => CallToolResult | Promise<CallToolResult>;
+export type MockedTools = Record<string, ToolResultGeneratorFn>;
+
+export class AccuracyTestingClient {
+    private mockedTools: MockedTools = {};
+    private recordedToolCalls: ToolCall[] = [];
+    private constructor(private readonly client: Awaited<ReturnType<typeof createMCPClient>>) {}
+
+    async close() {
+        await this.client?.close();
+    }
+
+    async vercelTools() {
+        const vercelTools = (await this.client?.tools()) ?? {};
+        const rewrappedVercelTools: typeof vercelTools = {};
+        for (const [toolName, tool] of Object.entries(vercelTools)) {
+            rewrappedVercelTools[toolName] = createVercelTool({
+                ...tool,
+                execute: async (args, options) => {
+                    this.recordedToolCalls.push({
+                        toolCallId: uuid(),
+                        toolName: toolName,
+                        parameters: args,
+                    });
+                    const toolResultGeneratorFn = this.mockedTools[toolName];
+                    if (toolResultGeneratorFn) {
+                        return await toolResultGeneratorFn(args);
+                    }
+
+                    return tool.execute(args, options);
+                },
+            });
+        }
+
+        return rewrappedVercelTools;
+    }
+
+    getToolCalls() {
+        return this.recordedToolCalls;
+    }
+
+    mockTools(mockedTools: MockedTools) {
+        this.mockedTools = mockedTools;
+    }
+
+    resetForTests() {
+        this.mockTools({});
+        this.recordedToolCalls = [];
+    }
+
+    static async initializeClient(mdbConnectionString: string) {
+        const clientTransport = new StdioClientTransport({
+            command: process.execPath,
+            args: [cliScriptPath, "--connectionString", mdbConnectionString],
+        });
+
+        const client = await createMCPClient({
+            transport: clientTransport,
+        });
+
+        return new AccuracyTestingClient(client);
+    }
+}

tests/accuracy/sdk/describe-accuracy-tests.ts

@@ -1,9 +1,8 @@
-import { Tool } from "@modelcontextprotocol/sdk/types.js";
-import { discoverMongoDBTools, TestTools, MockedTools } from "./test-tools.js";
 import { TestableModels } from "./models.js";
 import { ExpectedToolCall, parameterMatchingAccuracyScorer, toolCallingAccuracyScorer } from "./accuracy-scorers.js";
 import { Agent, getVercelToolCallingAgent } from "./agent.js";
-import { appendAccuracySnapshot } from "./accuracy-snapshot.js";
+import { prepareTestData, setupMongoDBIntegrationTest } from "../../integration/tools/mongodb/mongodbHelpers.js";
+import { AccuracyTestingClient, MockedTools } from "./accuracy-testing-client.js";
 
 export interface AccuracyTestConfig {
     systemPrompt?: string;
@@ -13,68 +12,71 @@ export interface AccuracyTestConfig {
     mockedTools: MockedTools;
 }
 
+export function describeSuite(suiteName: string, testConfigs: AccuracyTestConfig[]) {
+    return {
+        [suiteName]: testConfigs,
+    };
+}
+
 export function describeAccuracyTests(
-    suiteName: string,
     models: TestableModels,
-    accuracyTestConfigs: AccuracyTestConfig[]
+    accuracyTestConfigs: {
+        [suiteName: string]: AccuracyTestConfig[];
+    }
 ) {
-    const accuracyDatetime = process.env.MDB_ACCURACY_DATETIME;
-    const accuracyCommit = process.env.MDB_ACCURACY_COMMIT;
-
     if (!models.length) {
-        console.warn(`No models available to test ${suiteName}`);
-        return;
+        throw new Error("No models available to test!");
     }
 
     const eachModel = describe.each(models);
-    const eachTest = it.each(accuracyTestConfigs);
+    const eachSuite = describe.each(Object.keys(accuracyTestConfigs));
+
+    eachModel(`$modelName`, function (model) {
+        const mdbIntegration = setupMongoDBIntegrationTest();
+        const populateTestData = prepareTestData(mdbIntegration);
 
-    eachModel(`$modelName - ${suiteName}`, function (model) {
-        let mcpTools: Tool[];
-        let testTools: TestTools;
+        let testMCPClient: AccuracyTestingClient;
         let agent: Agent;
 
         beforeAll(async () => {
-            mcpTools = await discoverMongoDBTools();
+            testMCPClient = await AccuracyTestingClient.initializeClient(mdbIntegration.connectionString());
+            agent = getVercelToolCallingAgent();
         });
 
-        beforeEach(() => {
-            testTools = new TestTools(mcpTools);
-            agent = getVercelToolCallingAgent();
+        beforeEach(async () => {
+            await populateTestData();
+            testMCPClient.resetForTests();
+        });
+
+        afterAll(async () => {
+            await testMCPClient.close();
         });
 
-        eachTest("$prompt", async function (testConfig) {
-            testTools.mockTools(testConfig.mockedTools);
-            const toolsForModel = testTools.vercelAiTools();
-            const promptForModel = testConfig.injectConnectedAssumption
-                ? [testConfig.prompt, "(Assume that you are already connected to a MongoDB cluster!)"].join(" ")
-                : testConfig.prompt;
-            const conversation = await agent.prompt(promptForModel, model, toolsForModel);
-            const toolCalls = testTools.getToolCalls();
-            const toolCallingAccuracy = toolCallingAccuracyScorer(testConfig.expectedToolCalls, toolCalls);
-            const parameterMatchingAccuracy = parameterMatchingAccuracyScorer(testConfig.expectedToolCalls, toolCalls);
-            console.debug(`Conversation`, JSON.stringify(conversation, null, 2));
-            console.debug(`Tool calls`, JSON.stringify(toolCalls, null, 2));
-            console.debug(
-                "Tool calling accuracy: %s, Parameter Accuracy: %s",
-                toolCallingAccuracy,
-                parameterMatchingAccuracy
-            );
-            if (accuracyDatetime && accuracyCommit) {
-                await appendAccuracySnapshot({
-                    datetime: accuracyDatetime,
-                    commit: accuracyCommit,
-                    model: model.modelName,
-                    suite: suiteName,
-                    test: testConfig.prompt,
+        eachSuite("%s", function (suiteName) {
+            const eachTest = it.each(accuracyTestConfigs[suiteName] ?? []);
+
+            eachTest("$prompt", async function (testConfig) {
+                testMCPClient.mockTools(testConfig.mockedTools);
+                const toolsForModel = await testMCPClient.vercelTools();
+                const promptForModel = testConfig.injectConnectedAssumption
+                    ? [testConfig.prompt, "(Assume that you are already connected to a MongoDB cluster!)"].join(" ")
+                    : testConfig.prompt;
+                const conversation = await agent.prompt(promptForModel, model, toolsForModel);
+                const toolCalls = testMCPClient.getToolCalls();
+                const toolCallingAccuracy = toolCallingAccuracyScorer(testConfig.expectedToolCalls, toolCalls);
+                const parameterMatchingAccuracy = parameterMatchingAccuracyScorer(
+                    testConfig.expectedToolCalls,
+                    toolCalls
+                );
+                console.debug(testConfig.prompt);
+                console.debug(`Conversation`, JSON.stringify(conversation, null, 2));
+                // console.debug(`Tool calls`, JSON.stringify(toolCalls, null, 2));
+                console.debug(
+                    "Tool calling accuracy: %s, Parameter Accuracy: %s",
                     toolCallingAccuracy,
-                    parameterAccuracy: parameterMatchingAccuracy,
-                });
-            } else {
-                console.info(
-                    `Skipping accuracy snapshot update for ${model.modelName} - ${suiteName} - ${testConfig.prompt}`
+                    parameterMatchingAccuracy
                 );
-            }
+            });
         });
     });
 }