chore: add tests for metadata actions

src/tools/mongodb/metadata/collectionSchema.ts

@@ -1,7 +1,7 @@
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import { ToolArgs, OperationType } from "../../tool.js";
-import { parseSchema, SchemaField } from "mongodb-schema";
+import { getSimplifiedSchema } from "mongodb-schema";
 
 export class CollectionSchemaTool extends MongoDBToolBase {
     protected name = "collection-schema";
@@ -13,29 +13,31 @@ export class CollectionSchemaTool extends MongoDBToolBase {
     protected async execute({ database, collection }: ToolArgs<typeof DbOperationArgs>): Promise<CallToolResult> {
         const provider = await this.ensureConnected();
         const documents = await provider.find(database, collection, {}, { limit: 5 }).toArray();
-        const schema = await parseSchema(documents);
+        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 ${schema.fields.length} fields in the schema for \`${database}.${collection}\``,
+                    text: `Found ${fieldsCount} fields in the schema for "${database}.${collection}"`,
                     type: "text",
                 },
                 {
-                    text: this.formatFieldOutput(schema.fields),
+                    text: JSON.stringify(schema),
                     type: "text",
                 },
             ],
         };
     }
-
-    private formatFieldOutput(fields: SchemaField[]): string {
-        let result = "| Field | Type | Confidence |\n";
-        result += "|-------|------|-------------|\n";
-        for (const field of fields) {
-            const fieldType = Array.isArray(field.type) ? field.type.join(", ") : field.type;
-            result += `| ${field.name} | \`${fieldType}\` | ${(field.probability * 100).toFixed(0)}% |\n`;
-        }
-        return result;
-    }
 }

src/tools/mongodb/metadata/collectionStorageSize.ts

@@ -4,7 +4,7 @@ import { ToolArgs, OperationType } from "../../tool.js";
 
 export class CollectionStorageSizeTool extends MongoDBToolBase {
     protected name = "collection-storage-size";
-    protected description = "Gets the size of the collection in MB";
+    protected description = "Gets the size of the collection";
     protected argsShape = DbOperationArgs;
 
     protected operationType: OperationType = "metadata";
@@ -14,17 +14,55 @@ export class CollectionStorageSizeTool extends MongoDBToolBase {
         const [{ value }] = (await provider
             .aggregate(database, collection, [
                 { $collStats: { storageStats: {} } },
-                { $group: { _id: null, value: { $sum: "$storageStats.storageSize" } } },
+                { $group: { _id: null, value: { $sum: "$storageStats.size" } } },
             ])
             .toArray()) as [{ value: number }];
 
+        const { units, value: scaledValue } = CollectionStorageSizeTool.getStats(value);
+
         return {
             content: [
                 {
-                    text: `The size of \`${database}.${collection}\` is \`${(value / 1024 / 1024).toFixed(2)} MB\``,
+                    text: `The size of "${database}.${collection}" is \`${scaledValue.toFixed(2)} ${units}\``,
                     type: "text",
                 },
             ],
         };
     }
+
+    protected handleError(
+        error: unknown,
+        args: ToolArgs<typeof this.argsShape>
+    ): Promise<CallToolResult> | CallToolResult {
+        if (error instanceof Error && "codeName" in error && error.codeName === "NamespaceNotFound") {
+            return {
+                content: [
+                    {
+                        text: `The size of "${args.database}.${args.collection}" cannot be determined because the collection does not exist.`,
+                        type: "text",
+                    },
+                ],
+            };
+        }
+
+        return super.handleError(error, args);
+    }
+
+    private static getStats(value: number): { value: number; units: string } {
+        const kb = 1024;
+        const mb = kb * 1024;
+        const gb = mb * 1024;
+
+        if (value > gb) {
+            return { value: value / gb, units: "GB" };
+        }
+
+        if (value > mb) {
+            return { value: value / mb, units: "MB" };
+        }
+        if (value > kb) {
+            return { value: value / kb, units: "KB" };
+        }
+        return { value, units: "bytes" };
+    }
 }

src/tools/mongodb/mongodbTool.ts

@@ -1,5 +1,5 @@
 import { z } from "zod";
-import { ToolBase, ToolCategory } from "../tool.js";
+import { ToolArgs, ToolBase, ToolCategory } from "../tool.js";
 import { Session } from "../../session.js";
 import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
@@ -30,7 +30,10 @@ export abstract class MongoDBToolBase extends ToolBase {
         return this.session.serviceProvider;
     }
 
-    protected handleError(error: unknown): Promise<CallToolResult> | CallToolResult {
+    protected handleError(
+        error: unknown,
+        args: ToolArgs<typeof this.argsShape>
+    ): Promise<CallToolResult> | CallToolResult {
         if (error instanceof MongoDBError && error.code === ErrorCodes.NotConnectedToMongoDB) {
             return {
                 content: [
@@ -47,7 +50,7 @@ export abstract class MongoDBToolBase extends ToolBase {
             };
         }
 
-        return super.handleError(error);
+        return super.handleError(error, args);
     }
 
     protected async connectToMongoDB(connectionString: string): Promise<void> {

src/tools/tool.ts

@@ -44,7 +44,7 @@ export abstract class ToolBase {
             } catch (error: unknown) {
                 logger.error(mongoLogId(1_000_000), "tool", `Error executing ${this.name}: ${error as string}`);
 
-                return await this.handleError(error);
+                return await this.handleError(error, args[0] as ToolArgs<typeof this.argsShape>);
             }
         };
 
@@ -76,7 +76,11 @@ export abstract class ToolBase {
     }
 
     // This method is intended to be overridden by subclasses to handle errors
-    protected handleError(error: unknown): Promise<CallToolResult> | CallToolResult {
+    protected handleError(
+        error: unknown,
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        args: ToolArgs<typeof this.argsShape>
+    ): Promise<CallToolResult> | CallToolResult {
         return {
             content: [
                 {

tests/integration/helpers.ts

@@ -9,6 +9,7 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { MongoClient, ObjectId } from "mongodb";
 import { toIncludeAllMembers } from "jest-extended";
 import config from "../../src/config.js";
+import { McpError } from "@modelcontextprotocol/sdk/types.js";
 
 interface ParameterInfo {
     name: string;
@@ -223,10 +224,93 @@ export const dbOperationParameters: ParameterInfo[] = [
     { name: "collection", type: "string", description: "Collection name", required: true },
 ];
 
-export function validateParameters(tool: ToolInfo, parameters: ParameterInfo[]): void {
-    const toolParameters = getParameters(tool);
-    expect(toolParameters).toHaveLength(parameters.length);
-    expect(toolParameters).toIncludeAllMembers(parameters);
+export const dbOperationInvalidArgTests = [{}, { database: 123 }, { foo: "bar", database: "test" }, { database: [] }];
+
+export function validateToolMetadata(
+    integration: IntegrationTest,
+    name: string,
+    description: string,
+    parameters: ParameterInfo[]
+): void {
+    it("should have correct metadata", async () => {
+        const { tools } = await integration.mcpClient().listTools();
+        const tool = tools.find((tool) => tool.name === name)!;
+        expect(tool).toBeDefined();
+        expect(tool.description).toBe(description);
+
+        const toolParameters = getParameters(tool);
+        expect(toolParameters).toHaveLength(parameters.length);
+        expect(toolParameters).toIncludeAllMembers(parameters);
+    });
+}
+
+export function validateAutoConnectBehavior(
+    integration: IntegrationTest,
+    name: string,
+    validation: () => {
+        args: { [x: string]: unknown };
+        expectedResponse?: string;
+        validate?: (content: unknown) => void;
+    },
+    beforeEachImpl?: () => Promise<void>
+): void {
+    describe("when not connected", () => {
+        if (beforeEachImpl) {
+            beforeEach(() => beforeEachImpl());
+        }
+
+        it("connects automatically if connection string is configured", async () => {
+            config.connectionString = integration.connectionString();
+
+            const validationInfo = validation();
+
+            const response = await integration.mcpClient().callTool({
+                name,
+                arguments: validationInfo.args,
+            });
+
+            if (validationInfo.expectedResponse) {
+                const content = getResponseContent(response.content);
+                expect(content).toContain(validationInfo.expectedResponse);
+            }
+
+            if (validationInfo.validate) {
+                validationInfo.validate(response.content);
+            }
+        });
+
+        it("throws an error if connection string is not configured", async () => {
+            const response = await integration.mcpClient().callTool({
+                name,
+                arguments: validation().args,
+            });
+            const content = getResponseContent(response.content);
+            expect(content).toContain("You need to connect to a MongoDB instance before you can access its data.");
+        });
+    });
+}
+
+export function validateThrowsForInvalidArguments(
+    integration: IntegrationTest,
+    name: string,
+    args: { [x: string]: unknown }[]
+): void {
+    describe("with invalid arguments", () => {
+        for (const arg of args) {
+            it(`throws a schema error for: ${JSON.stringify(arg)}`, async () => {
+                await integration.connectMcpClient();
+                try {
+                    await integration.mcpClient().callTool({ name, arguments: arg });
+                    expect.fail("Expected an error to be thrown");
+                } catch (error) {
+                    expect(error).toBeInstanceOf(McpError);
+                    const mcpError = error as McpError;
+                    expect(mcpError.code).toEqual(-32602);
+                    expect(mcpError.message).toContain(`Invalid arguments for tool ${name}`);
+                }
+            });
+        }
+    });
 }
 
 export function describeAtlas(name: number | string | Function | jest.FunctionLike, fn: jest.EmptyFunction) {

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

@@ -1,50 +1,24 @@
 import {
     getResponseContent,
-    validateParameters,
     dbOperationParameters,
     setupIntegrationTest,
+    validateToolMetadata,
+    validateAutoConnectBehavior,
+    validateThrowsForInvalidArguments,
+    dbOperationInvalidArgTests,
 } from "../../../helpers.js";
-import { toIncludeSameMembers } from "jest-extended";
-import { McpError } from "@modelcontextprotocol/sdk/types.js";
-import { ObjectId } from "bson";
-import config from "../../../../../src/config.js";
 
 describe("createCollection tool", () => {
     const integration = setupIntegrationTest();
 
-    it("should have correct metadata", async () => {
-        const { tools } = await integration.mcpClient().listTools();
-        const listCollections = tools.find((tool) => tool.name === "create-collection")!;
-        expect(listCollections).toBeDefined();
-        expect(listCollections.description).toBe(
-            "Creates a new collection in a database. If the database doesn't exist, it will be created automatically."
-        );
+    validateToolMetadata(
+        integration,
+        "create-collection",
+        "Creates a new collection in a database. If the database doesn't exist, it will be created automatically.",
+        dbOperationParameters
+    );
 
-        validateParameters(listCollections, dbOperationParameters);
-    });
-
-    describe("with invalid arguments", () => {
-        const args = [
-            {},
-            { database: 123, collection: "bar" },
-            { foo: "bar", database: "test", collection: "bar" },
-            { collection: [], database: "test" },
-        ];
-        for (const arg of args) {
-            it(`throws a schema error for: ${JSON.stringify(arg)}`, async () => {
-                await integration.connectMcpClient();
-                try {
-                    await integration.mcpClient().callTool({ name: "create-collection", arguments: arg });
-                    expect.fail("Expected an error to be thrown");
-                } catch (error) {
-                    expect(error).toBeInstanceOf(McpError);
-                    const mcpError = error as McpError;
-                    expect(mcpError.code).toEqual(-32602);
-                    expect(mcpError.message).toContain("Invalid arguments for tool create-collection");
-                }
-            });
-        }
-    });
+    validateThrowsForInvalidArguments(integration, "create-collection", dbOperationInvalidArgTests);
 
     describe("with non-existent database", () => {
         it("creates a new collection", async () => {
@@ -114,25 +88,10 @@ describe("createCollection tool", () => {
         });
     });
 
-    describe("when not connected", () => {
-        it("connects automatically if connection string is configured", async () => {
-            config.connectionString = integration.connectionString();
-
-            const response = await integration.mcpClient().callTool({
-                name: "create-collection",
-                arguments: { database: integration.randomDbName(), collection: "new-collection" },
-            });
-            const content = getResponseContent(response.content);
-            expect(content).toEqual(`Collection "new-collection" created in database "${integration.randomDbName()}".`);
-        });
-
-        it("throws an error if connection string is not configured", async () => {
-            const response = await integration.mcpClient().callTool({
-                name: "create-collection",
-                arguments: { database: integration.randomDbName(), collection: "new-collection" },
-            });
-            const content = getResponseContent(response.content);
-            expect(content).toContain("You need to connect to a MongoDB instance before you can access its data.");
-        });
+    validateAutoConnectBehavior(integration, "create-collection", () => {
+        return {
+            args: { database: integration.randomDbName(), collection: "new-collection" },
+            expectedResponse: `Collection "new-collection" created in database "${integration.randomDbName()}".`,
+        };
     });
 });