chore: add tests for collection-storage-size

nirinchev · nirinchev · commit 2c736b9629ce · 2025-04-23T01:13:50.000+02:00
diff --git a/src/tools/mongodb/metadata/collectionStorageSize.ts b/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" };
+    }
 }
diff --git a/src/tools/mongodb/mongodbTool.ts b/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> {
diff --git a/src/tools/tool.ts b/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: [
                 {
diff --git a/tests/integration/helpers.ts b/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;
@@ -210,6 +211,8 @@ export const dbOperationParameters: ParameterInfo[] = [
     { name: "collection", type: "string", description: "Collection name", required: true },
 ];
 
+export const dbOperationInvalidArgTests = [{}, { database: 123 }, { foo: "bar", database: "test" }, { database: [] }];
+
 export async function validateToolMetadata(
     mcpClient: Client,
     name: string,
@@ -264,3 +267,24 @@ export function validateAutoConnectBehavior(
         expect(content).toContain("You need to connect to a MongoDB instance before you can access its data.");
     });
 }
+
+export function validateThrowsForInvalidArguments(
+    integration: IntegrationTestSetup,
+    name: string,
+    args: { [x: string]: unknown }[]
+): void {
+    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}`);
+            }
+        });
+    }
+}
diff --git a/tests/integration/tools/mongodb/create/createCollection.test.ts b/tests/integration/tools/mongodb/create/createCollection.test.ts
@@ -4,6 +4,8 @@ import {
     setupIntegrationTest,
     validateToolMetadata,
     validateAutoConnectBehavior,
+    validateThrowsForInvalidArguments,
+    dbOperationInvalidArgTests,
 } from "../../../helpers.js";
 import { toIncludeSameMembers } from "jest-extended";
 import { McpError } from "@modelcontextprotocol/sdk/types.js";
@@ -21,26 +23,7 @@ describe("createCollection tool", () => {
     });
 
     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", () => {
diff --git a/tests/integration/tools/mongodb/create/createIndex.test.ts b/tests/integration/tools/mongodb/create/createIndex.test.ts
@@ -4,6 +4,7 @@ import {
     setupIntegrationTest,
     validateToolMetadata,
     validateAutoConnectBehavior,
+    validateThrowsForInvalidArguments,
 } from "../../../helpers.js";
 import { McpError } from "@modelcontextprotocol/sdk/types.js";
 import { IndexDirection } from "mongodb";
@@ -31,28 +32,14 @@ describe("createIndex tool", () => {
     });
 
     describe("with invalid arguments", () => {
-        const args = [
+        validateThrowsForInvalidArguments(integration, "create-index", [
             {},
             { collection: "bar", database: 123, keys: { foo: 1 } },
             { collection: "bar", database: "test", keys: { foo: 5 } },
             { collection: [], database: "test", keys: { foo: 1 } },
             { collection: "bar", database: "test", keys: { foo: 1 }, name: 123 },
             { collection: "bar", database: "test", keys: "foo", name: "my-index" },
-        ];
-        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-index", 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-index");
-                }
-            });
-        }
+        ]);
     });
 
     const validateIndex = async (collection: string, expected: { name: string; key: object }[]) => {
diff --git a/tests/integration/tools/mongodb/create/insertMany.test.ts b/tests/integration/tools/mongodb/create/insertMany.test.ts
@@ -4,6 +4,7 @@ import {
     setupIntegrationTest,
     validateToolMetadata,
     validateAutoConnectBehavior,
+    validateThrowsForInvalidArguments,
 } from "../../../helpers.js";
 import { McpError } from "@modelcontextprotocol/sdk/types.js";
 import config from "../../../../../src/config.js";
@@ -30,27 +31,13 @@ describe("insertMany tool", () => {
     });
 
     describe("with invalid arguments", () => {
-        const args = [
+        validateThrowsForInvalidArguments(integration, "insert-many", [
             {},
             { collection: "bar", database: 123, documents: [] },
             { collection: [], database: "test", documents: [] },
             { collection: "bar", database: "test", documents: "my-document" },
             { collection: "bar", database: "test", documents: { name: "Peter" } },
-        ];
-        for (const arg of args) {
-            it(`throws a schema error for: ${JSON.stringify(arg)}`, async () => {
-                await integration.connectMcpClient();
-                try {
-                    await integration.mcpClient().callTool({ name: "insert-many", 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 insert-many");
-                }
-            });
-        }
+        ]);
     });
 
     const validateDocuments = async (collection: string, expectedDocuments: object[]) => {
diff --git a/tests/integration/tools/mongodb/delete/deleteMany.test.ts b/tests/integration/tools/mongodb/delete/deleteMany.test.ts
@@ -4,9 +4,8 @@ import {
     setupIntegrationTest,
     validateToolMetadata,
     validateAutoConnectBehavior,
+    validateThrowsForInvalidArguments,
 } from "../../../helpers.js";
-import { McpError } from "@modelcontextprotocol/sdk/types.js";
-import config from "../../../../../src/config.js";
 
 describe("deleteMany tool", () => {
     const integration = setupIntegrationTest();
@@ -30,27 +29,13 @@ describe("deleteMany tool", () => {
     });
 
     describe("with invalid arguments", () => {
-        const args = [
+        validateThrowsForInvalidArguments(integration, "delete-many", [
             {},
             { collection: "bar", database: 123, filter: {} },
             { collection: [], database: "test", filter: {} },
             { collection: "bar", database: "test", filter: "my-document" },
             { collection: "bar", database: "test", filter: [{ name: "Peter" }] },
-        ];
-        for (const arg of args) {
-            it(`throws a schema error for: ${JSON.stringify(arg)}`, async () => {
-                await integration.connectMcpClient();
-                try {
-                    await integration.mcpClient().callTool({ name: "delete-many", 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 delete-many");
-                }
-            });
-        }
+        ]);
     });
 
     it("doesn't create the collection if it doesn't exist", async () => {
diff --git a/tests/integration/tools/mongodb/delete/dropCollection.test.ts b/tests/integration/tools/mongodb/delete/dropCollection.test.ts
@@ -4,9 +4,9 @@ import {
     setupIntegrationTest,
     validateToolMetadata,
     validateAutoConnectBehavior,
+    validateThrowsForInvalidArguments,
+    dbOperationInvalidArgTests,
 } from "../../../helpers.js";
-import { McpError } from "@modelcontextprotocol/sdk/types.js";
-import config from "../../../../../src/config.js";
 
 describe("dropCollection tool", () => {
     const integration = setupIntegrationTest();
@@ -21,26 +21,7 @@ describe("dropCollection tool", () => {
     });
 
     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: "drop-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 drop-collection");
-                }
-            });
-        }
+        validateThrowsForInvalidArguments(integration, "drop-collection", dbOperationInvalidArgTests);
     });
 
     it("can drop non-existing collection", async () => {
diff --git a/tests/integration/tools/mongodb/delete/dropDatabase.test.ts b/tests/integration/tools/mongodb/delete/dropDatabase.test.ts
@@ -4,9 +4,9 @@ import {
     setupIntegrationTest,
     validateToolMetadata,
     validateAutoConnectBehavior,
+    validateThrowsForInvalidArguments,
+    dbOperationInvalidArgTests,
 } from "../../../helpers.js";
-import { McpError } from "@modelcontextprotocol/sdk/types.js";
-import config from "../../../../../src/config.js";
 
 describe("dropDatabase tool", () => {
     const integration = setupIntegrationTest();
@@ -21,21 +21,7 @@ describe("dropDatabase tool", () => {
     });
 
     describe("with invalid arguments", () => {
-        const args = [{}, { database: 123 }, { foo: "bar", 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: "drop-database", 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 drop-database");
-                }
-            });
-        }
+        validateThrowsForInvalidArguments(integration, "drop-database", dbOperationInvalidArgTests);
     });
 
     it("can drop non-existing database", async () => {
diff --git a/tests/integration/tools/mongodb/metadata/collectionSchema.test.ts b/tests/integration/tools/mongodb/metadata/collectionSchema.test.ts
diff --git a/tests/integration/tools/mongodb/metadata/collectionStorageSize.test.ts b/tests/integration/tools/mongodb/metadata/collectionStorageSize.test.ts
diff --git a/tests/integration/tools/mongodb/metadata/listCollections.test.ts b/tests/integration/tools/mongodb/metadata/listCollections.test.ts
diff --git a/tests/integration/tools/mongodb/read/count.test.ts b/tests/integration/tools/mongodb/read/count.test.ts

Original file line number	Diff line number	Diff line change
`@@ -44,7 +44,7 @@ export abstract class ToolBase {`
`44`	`44`	`} catch (error: unknown) {`
`45`	`45`	logger.error(mongoLogId(1_000_000), "tool", `Error executing ${this.name}: ${error as string}`);
`46`	`46`
`47`		`- return await this.handleError(error);`
	`47`	`+ return await this.handleError(error, args[0] as ToolArgs<typeof this.argsShape>);`
`48`	`48`	`}`
`49`	`49`	`};`
`50`	`50`
`@@ -76,7 +76,11 @@ export abstract class ToolBase {`
`76`	`76`	`}`
`77`	`77`
`78`	`78`	`// This method is intended to be overridden by subclasses to handle errors`
`79`		`- protected handleError(error: unknown): Promise<CallToolResult> \| CallToolResult {`
	`79`	`+ protected handleError(`
	`80`	`+ error: unknown,`
	`81`	`+ // eslint-disable-next-line @typescript-eslint/no-unused-vars`
	`82`	`+ args: ToolArgs<typeof this.argsShape>`
	`83`	`+ ): Promise<CallToolResult> \| CallToolResult {`
`80`	`84`	`return {`
`81`	`85`	`content: [`
`82`	`86`	`{`