chore: add tests for db-stats (#99)

nirinchev · web-flow · commit 673fca7faceb · 2025-04-24T13:32:38.000+02:00
diff --git a/src/tools/mongodb/metadata/dbStats.ts b/src/tools/mongodb/metadata/dbStats.ts
@@ -1,6 +1,7 @@
 import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
 import { ToolArgs, OperationType } from "../../tool.js";
+import { EJSON } from "bson";
 
 export class DbStatsTool extends MongoDBToolBase {
     protected name = "db-stats";
@@ -21,7 +22,11 @@ export class DbStatsTool extends MongoDBToolBase {
         return {
             content: [
                 {
-                    text: `Statistics for database ${database}: ${JSON.stringify(result)}`,
+                    text: `Statistics for database ${database}`,
+                    type: "text",
+                },
+                {
+                    text: EJSON.stringify(result),
                     type: "text",
                 },
             ],
diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts
@@ -149,12 +149,27 @@ export function getParameters(tool: ToolInfo): ParameterInfo[] {
         });
 }
 
-export const dbOperationParameters: ParameterInfo[] = [
+export const databaseParameters: ParameterInfo[] = [
     { name: "database", type: "string", description: "Database name", required: true },
+];
+
+export const databaseCollectionParameters: ParameterInfo[] = [
+    ...databaseParameters,
     { name: "collection", type: "string", description: "Collection name", required: true },
 ];
 
-export const dbOperationInvalidArgTests = [{}, { database: 123 }, { foo: "bar", database: "test" }, { database: [] }];
+export const databaseCollectionInvalidArgs = [
+    {},
+    { database: "test" },
+    { collection: "foo" },
+    { database: 123, collection: "foo" },
+    { database: "test", collection: "foo", extra: "bar" },
+    { database: "test", collection: 123 },
+    { database: [], collection: "foo" },
+    { database: "test", collection: [] },
+];
+
+export const databaseInvalidArgs = [{}, { database: 123 }, { database: [] }, { database: "test", extra: "bar" }];
 
 export function validateToolMetadata(
     integration: IntegrationTest,
diff --git a/tests/integration/tools/mongodb/create/createCollection.test.ts b/tests/integration/tools/mongodb/create/createCollection.test.ts
@@ -2,21 +2,21 @@ import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelp
 
 import {
     getResponseContent,
-    dbOperationParameters,
+    databaseCollectionParameters,
     validateToolMetadata,
     validateThrowsForInvalidArguments,
-    dbOperationInvalidArgTests,
+    databaseCollectionInvalidArgs,
 } from "../../../helpers.js";
 
 describeWithMongoDB("createCollection tool", (integration) => {
     validateToolMetadata(
         integration,
         "create-collection",
         "Creates a new collection in a database. If the database doesn't exist, it will be created automatically.",
-        dbOperationParameters
+        databaseCollectionParameters
     );
 
-    validateThrowsForInvalidArguments(integration, "create-collection", dbOperationInvalidArgTests);
+    validateThrowsForInvalidArguments(integration, "create-collection", databaseCollectionInvalidArgs);
 
     describe("with non-existent database", () => {
         it("creates a new collection", async () => {
diff --git a/tests/integration/tools/mongodb/create/createIndex.test.ts b/tests/integration/tools/mongodb/create/createIndex.test.ts
@@ -2,15 +2,15 @@ import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelp
 
 import {
     getResponseContent,
-    dbOperationParameters,
+    databaseCollectionParameters,
     validateToolMetadata,
     validateThrowsForInvalidArguments,
 } from "../../../helpers.js";
 import { IndexDirection } from "mongodb";
 
 describeWithMongoDB("createIndex tool", (integration) => {
     validateToolMetadata(integration, "create-index", "Create an index for a collection", [
-        ...dbOperationParameters,
+        ...databaseCollectionParameters,
         {
             name: "keys",
             type: "object",
diff --git a/tests/integration/tools/mongodb/create/insertMany.test.ts b/tests/integration/tools/mongodb/create/insertMany.test.ts
@@ -2,14 +2,14 @@ import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelp
 
 import {
     getResponseContent,
-    dbOperationParameters,
+    databaseCollectionParameters,
     validateToolMetadata,
     validateThrowsForInvalidArguments,
 } from "../../../helpers.js";
 
 describeWithMongoDB("insertMany tool", (integration) => {
     validateToolMetadata(integration, "insert-many", "Insert an array of documents into a MongoDB collection", [
-        ...dbOperationParameters,
+        ...databaseCollectionParameters,
         {
             name: "documents",
             type: "array",
diff --git a/tests/integration/tools/mongodb/delete/deleteMany.test.ts b/tests/integration/tools/mongodb/delete/deleteMany.test.ts
@@ -2,7 +2,7 @@ import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelp
 
 import {
     getResponseContent,
-    dbOperationParameters,
+    databaseCollectionParameters,
     validateToolMetadata,
     validateThrowsForInvalidArguments,
 } from "../../../helpers.js";
@@ -13,7 +13,7 @@ describeWithMongoDB("deleteMany tool", (integration) => {
         "delete-many",
         "Removes all documents that match the filter from a MongoDB collection",
         [
-            ...dbOperationParameters,
+            ...databaseCollectionParameters,
             {
                 name: "filter",
                 type: "object",
diff --git a/tests/integration/tools/mongodb/delete/dropCollection.test.ts b/tests/integration/tools/mongodb/delete/dropCollection.test.ts
@@ -2,21 +2,21 @@ import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelp
 
 import {
     getResponseContent,
-    dbOperationParameters,
+    databaseCollectionParameters,
     validateToolMetadata,
     validateThrowsForInvalidArguments,
-    dbOperationInvalidArgTests,
+    databaseCollectionInvalidArgs,
 } from "../../../helpers.js";
 
 describeWithMongoDB("dropCollection tool", (integration) => {
     validateToolMetadata(
         integration,
         "drop-collection",
         "Removes a collection or view from the database. The method also removes any indexes associated with the dropped collection.",
-        dbOperationParameters
+        databaseCollectionParameters
     );
 
-    validateThrowsForInvalidArguments(integration, "drop-collection", dbOperationInvalidArgTests);
+    validateThrowsForInvalidArguments(integration, "drop-collection", databaseCollectionInvalidArgs);
 
     it("can drop non-existing collection", async () => {
         await integration.connectMcpClient();
diff --git a/tests/integration/tools/mongodb/delete/dropDatabase.test.ts b/tests/integration/tools/mongodb/delete/dropDatabase.test.ts
@@ -2,21 +2,21 @@ import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelp
 
 import {
     getResponseContent,
-    dbOperationParameters,
     validateToolMetadata,
     validateThrowsForInvalidArguments,
-    dbOperationInvalidArgTests,
+    databaseParameters,
+    databaseInvalidArgs,
 } from "../../../helpers.js";
 
 describeWithMongoDB("dropDatabase tool", (integration) => {
     validateToolMetadata(
         integration,
         "drop-database",
         "Removes the specified database, deleting the associated data files",
-        [dbOperationParameters.find((d) => d.name === "database")!]
+        databaseParameters
     );
 
-    validateThrowsForInvalidArguments(integration, "drop-database", dbOperationInvalidArgTests);
+    validateThrowsForInvalidArguments(integration, "drop-database", databaseInvalidArgs);
 
     it("can drop non-existing database", async () => {
         let { databases } = await integration.mongoClient().db("").admin().listDatabases();
diff --git a/tests/integration/tools/mongodb/metadata/collectionSchema.test.ts b/tests/integration/tools/mongodb/metadata/collectionSchema.test.ts
@@ -3,10 +3,10 @@ import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelp
 import {
     getResponseElements,
     getResponseContent,
-    dbOperationParameters,
+    databaseCollectionParameters,
     validateToolMetadata,
     validateThrowsForInvalidArguments,
-    dbOperationInvalidArgTests,
+    databaseCollectionInvalidArgs,
 } from "../../../helpers.js";
 import { Document } from "bson";
 import { OptionalId } from "mongodb";
@@ -17,10 +17,10 @@ describeWithMongoDB("collectionSchema tool", (integration) => {
         integration,
         "collection-schema",
         "Describe the schema for a collection",
-        dbOperationParameters
+        databaseCollectionParameters
     );
 
-    validateThrowsForInvalidArguments(integration, "collection-schema", dbOperationInvalidArgTests);
+    validateThrowsForInvalidArguments(integration, "collection-schema", databaseCollectionInvalidArgs);
 
     describe("with non-existent database", () => {
         it("returns empty schema", async () => {
diff --git a/tests/integration/tools/mongodb/metadata/collectionStorageSize.test.ts b/tests/integration/tools/mongodb/metadata/collectionStorageSize.test.ts
@@ -2,9 +2,9 @@ import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelp
 
 import {
     getResponseContent,
-    dbOperationParameters,
+    databaseCollectionParameters,
+    databaseCollectionInvalidArgs,
     validateToolMetadata,
-    dbOperationInvalidArgTests,
     validateThrowsForInvalidArguments,
 } from "../../../helpers.js";
 import * as crypto from "crypto";
@@ -14,13 +14,13 @@ describeWithMongoDB("collectionStorageSize tool", (integration) => {
         integration,
         "collection-storage-size",
         "Gets the size of the collection",
-        dbOperationParameters
+        databaseCollectionParameters
     );
 
-    validateThrowsForInvalidArguments(integration, "collection-storage-size", dbOperationInvalidArgTests);
+    validateThrowsForInvalidArguments(integration, "collection-storage-size", databaseCollectionInvalidArgs);
 
     describe("with non-existent database", () => {
-        it("returns 0 MB", async () => {
+        it("returns an error", async () => {
             await integration.connectMcpClient();
             const response = await integration.mcpClient().callTool({
                 name: "collection-storage-size",
diff --git a/tests/integration/tools/mongodb/metadata/dbStats.test.ts b/tests/integration/tools/mongodb/metadata/dbStats.test.ts
@@ -0,0 +1,96 @@
+import { ObjectId } from "bson";
+import {
+    databaseParameters,
+    validateToolMetadata,
+    validateThrowsForInvalidArguments,
+    databaseInvalidArgs,
+    getResponseElements,
+} from "../../../helpers.js";
+import * as crypto from "crypto";
+import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelpers.js";
+
+describeWithMongoDB("dbStats tool", (integration) => {
+    validateToolMetadata(
+        integration,
+        "db-stats",
+        "Returns statistics that reflect the use state of a single database",
+        databaseParameters
+    );
+
+    validateThrowsForInvalidArguments(integration, "db-stats", databaseInvalidArgs);
+
+    describe("with non-existent database", () => {
+        it("returns an error", async () => {
+            await integration.connectMcpClient();
+            const response = await integration.mcpClient().callTool({
+                name: "db-stats",
+                arguments: { database: integration.randomDbName() },
+            });
+            const elements = getResponseElements(response.content);
+            expect(elements).toHaveLength(2);
+            expect(elements[0].text).toBe(`Statistics for database ${integration.randomDbName()}`);
+
+            const stats = JSON.parse(elements[1].text);
+            expect(stats.db).toBe(integration.randomDbName());
+            expect(stats.collections).toBe(0);
+            expect(stats.storageSize).toBe(0);
+        });
+    });
+
+    describe("with existing database", () => {
+        const testCases = [
+            {
+                collections: {
+                    foos: 3,
+                },
+                name: "single collection",
+            },
+            {
+                collections: {
+                    foos: 2,
+                    bars: 5,
+                },
+                name: "multiple collections",
+            },
+        ];
+        for (const test of testCases) {
+            it(`returns correct stats for ${test.name}`, async () => {
+                for (const [name, count] of Object.entries(test.collections)) {
+                    const objects = Array(count)
+                        .fill(0)
+                        .map(() => {
+                            return { data: crypto.randomBytes(1024), _id: new ObjectId() };
+                        });
+                    await integration.mongoClient().db(integration.randomDbName()).collection(name).insertMany(objects);
+                }
+
+                await integration.connectMcpClient();
+                const response = await integration.mcpClient().callTool({
+                    name: "db-stats",
+                    arguments: { database: integration.randomDbName() },
+                });
+                const elements = getResponseElements(response.content);
+                expect(elements).toHaveLength(2);
+                expect(elements[0].text).toBe(`Statistics for database ${integration.randomDbName()}`);
+
+                const stats = JSON.parse(elements[1].text);
+                expect(stats.db).toBe(integration.randomDbName());
+                expect(stats.collections).toBe(Object.entries(test.collections).length);
+                expect(stats.storageSize).toBeGreaterThan(1024);
+                expect(stats.objects).toBe(Object.values(test.collections).reduce((a, b) => a + b, 0));
+            });
+        }
+    });
+
+    describe("when not connected", () => {
+        validateAutoConnectBehavior(integration, "db-stats", () => {
+            return {
+                args: {
+                    database: integration.randomDbName(),
+                    collection: "foo",
+                },
+                expectedResponse: `Statistics for database ${integration.randomDbName()}`,
+            };
+        });
+    });
+});
diff --git a/tests/integration/tools/mongodb/metadata/listCollections.test.ts b/tests/integration/tools/mongodb/metadata/listCollections.test.ts
@@ -5,15 +5,19 @@ import {
     getResponseContent,
     validateToolMetadata,
     validateThrowsForInvalidArguments,
-    dbOperationInvalidArgTests,
+    databaseInvalidArgs,
+    databaseParameters,
 } from "../../../helpers.js";
 
 describeWithMongoDB("listCollections tool", (integration) => {
-    validateToolMetadata(integration, "list-collections", "List all collections for a given database", [
-        { name: "database", description: "Database name", type: "string", required: true },
-    ]);
+    validateToolMetadata(
+        integration,
+        "list-collections",
+        "List all collections for a given database",
+        databaseParameters
+    );
 
-    validateThrowsForInvalidArguments(integration, "list-collections", dbOperationInvalidArgTests);
+    validateThrowsForInvalidArguments(integration, "list-collections", databaseInvalidArgs);
 
     describe("with non-existent database", () => {
         it("returns no collections", async () => {
diff --git a/tests/integration/tools/mongodb/metadata/listDatabases.test.ts b/tests/integration/tools/mongodb/metadata/listDatabases.test.ts
@@ -1,5 +1,4 @@
 import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelpers.js";
-
 import { getResponseElements, getParameters } from "../../../helpers.js";
 
 describeWithMongoDB("listDatabases tool", (integration) => {
@@ -21,7 +20,7 @@ describeWithMongoDB("listDatabases tool", (integration) => {
             const response = await integration.mcpClient().callTool({ name: "list-databases", arguments: {} });
             const dbNames = getDbNames(response.content);
 
-            expect(defaultDatabases).toIncludeAllMembers(defaultDatabases);
+            expect(defaultDatabases).toIncludeAllMembers(dbNames);
         });
     });
 
diff --git a/tests/integration/tools/mongodb/read/count.test.ts b/tests/integration/tools/mongodb/read/count.test.ts
@@ -2,7 +2,7 @@ import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelp
 
 import {
     getResponseContent,
-    dbOperationParameters,
+    databaseCollectionParameters,
     validateToolMetadata,
     validateThrowsForInvalidArguments,
 } from "../../../helpers.js";
@@ -16,7 +16,7 @@ describeWithMongoDB("count tool", (integration) => {
             type: "object",
             required: false,
         },
-        ...dbOperationParameters,
+        ...databaseCollectionParameters,
     ]);
 
     validateThrowsForInvalidArguments(integration, "count", [
