chore: tests for collection-schema tool

Author: himanshusinghs

src/tools/mongodb/metadata/collectionSchema.ts

@@ -1,7 +1,38 @@
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import { ToolArgs, OperationType } from "../../tool.js";
-import { getSimplifiedSchema } from "mongodb-schema";
+import { getSimplifiedSchema, SimplifiedSchema } from "mongodb-schema";
+
+export function collectionSchemaResponse(
+    database: string,
+    collection: string,
+    schema: SimplifiedSchema
+): CallToolResult {
+    const fieldsCount = Object.entries(schema).length;
+    if (fieldsCount === 0) {
+        return {
+            content: [
+                {
+                    text: `Could not deduce the schema for "${database}.${collection}". This may be because it doesn't exist or is empty.`,
+                    type: "text",
+                },
+            ],
+        };
+    }
+
+    return {
+        content: [
+            {
+                text: `Found ${fieldsCount} fields in the schema for "${database}.${collection}"`,
+                type: "text",
+            },
+            {
+                text: JSON.stringify(schema),
+                type: "text",
+            },
+        ],
+    };
+}
 
 export class CollectionSchemaTool extends MongoDBToolBase {
     protected name = "collection-schema";
@@ -14,30 +45,6 @@ export class CollectionSchemaTool extends MongoDBToolBase {
         const provider = await this.ensureConnected();
         const documents = await provider.find(database, collection, {}, { limit: 5 }).toArray();
         const schema = await getSimplifiedSchema(documents);
-
-        const fieldsCount = Object.entries(schema).length;
-        if (fieldsCount === 0) {
-            return {
-                content: [
-                    {
-                        text: `Could not deduce the schema for "${database}.${collection}". This may be because it doesn't exist or is empty.`,
-                        type: "text",
-                    },
-                ],
-            };
-        }
-
-        return {
-            content: [
-                {
-                    text: `Found ${fieldsCount} fields in the schema for "${database}.${collection}"`,
-                    type: "text",
-                },
-                {
-                    text: JSON.stringify(schema),
-                    type: "text",
-                },
-            ],
-        };
+        return collectionSchemaResponse(database, collection, schema);
     }
 }

tests/accuracy/collection-schema.test.ts

@@ -0,0 +1,47 @@
+import { describeAccuracyTests } from "./sdk/describe-accuracy-tests.js";
+import { getAvailableModels } from "./sdk/models.js";
+import { AccuracyTestConfig } from "./sdk/describe-accuracy-tests.js";
+import { collectionSchemaResponse } from "../../src/tools/mongodb/metadata/collectionSchema.js";
+import { getSimplifiedSchema } from "mongodb-schema";
+
+function callsCollectionSchema(prompt: string): AccuracyTestConfig {
+    return {
+        injectConnectedAssumption: true,
+        prompt: prompt,
+        mockedTools: {
+            "collection-schema": async function collectionSchema() {
+                return collectionSchemaResponse(
+                    "db1",
+                    "coll1",
+                    await getSimplifiedSchema([
+                        {
+                            name: "Sample name1",
+                            dob: "28.11.2001",
+                            location: "NY",
+                        },
+                        {
+                            name: "Sample name1",
+                            dob: "28.11.2001",
+                            location: "NY",
+                            title: "Dr.",
+                        },
+                    ])
+                );
+            },
+        },
+        expectedToolCalls: [
+            {
+                toolName: "collection-schema",
+                parameters: {
+                    database: "db1",
+                    collection: "coll1",
+                },
+            },
+        ],
+    };
+}
+
+describeAccuracyTests("collection-schema", getAvailableModels(), [
+    callsCollectionSchema("Is there a title field in 'db1.coll1' namespace?"),
+    callsCollectionSchema("What is the type of value stored in title field in coll1 collection in db1 database?"),
+]);