chore: add tests for db-stats

nirinchev · nirinchev · commit b959b6a34911 · 2025-04-23T15:47:13.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
@@ -222,12 +222,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
@@ -1,11 +1,11 @@
 import {
     getResponseContent,
-    dbOperationParameters,
+    databaseCollectionParameters,
     setupIntegrationTest,
     validateToolMetadata,
     validateAutoConnectBehavior,
     validateThrowsForInvalidArguments,
-    dbOperationInvalidArgTests,
+    databaseCollectionInvalidArgs,
 } from "../../../helpers.js";
 
 describe("createCollection tool", () => {
@@ -15,10 +15,10 @@ describe("createCollection tool", () => {
         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
@@ -1,6 +1,6 @@
 import {
     getResponseContent,
-    dbOperationParameters,
+    databaseCollectionParameters,
     setupIntegrationTest,
     validateToolMetadata,
     validateAutoConnectBehavior,
@@ -12,7 +12,7 @@ describe("createIndex tool", () => {
     const integration = setupIntegrationTest();
 
     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
@@ -1,6 +1,6 @@
 import {
     getResponseContent,
-    dbOperationParameters,
+    databaseCollectionParameters,
     setupIntegrationTest,
     validateToolMetadata,
     validateAutoConnectBehavior,
@@ -11,7 +11,7 @@ describe("insertMany tool", () => {
     const integration = setupIntegrationTest();
 
     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
@@ -1,6 +1,6 @@
 import {
     getResponseContent,
-    dbOperationParameters,
+    databaseCollectionParameters,
     setupIntegrationTest,
     validateToolMetadata,
     validateAutoConnectBehavior,
@@ -15,7 +15,7 @@ describe("deleteMany tool", () => {
         "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
@@ -1,11 +1,11 @@
 import {
     getResponseContent,
-    dbOperationParameters,
+    databaseCollectionParameters,
     setupIntegrationTest,
     validateToolMetadata,
     validateAutoConnectBehavior,
     validateThrowsForInvalidArguments,
-    dbOperationInvalidArgTests,
+    databaseCollectionInvalidArgs,
 } from "../../../helpers.js";
 
 describe("dropCollection tool", () => {
@@ -15,10 +15,10 @@ describe("dropCollection tool", () => {
         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
@@ -1,11 +1,11 @@
 import {
     getResponseContent,
-    dbOperationParameters,
     setupIntegrationTest,
     validateToolMetadata,
     validateAutoConnectBehavior,
     validateThrowsForInvalidArguments,
-    dbOperationInvalidArgTests,
+    databaseParameters,
+    databaseInvalidArgs,
 } from "../../../helpers.js";
 
 describe("dropDatabase tool", () => {
@@ -15,10 +15,10 @@ describe("dropDatabase tool", () => {
         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
@@ -2,11 +2,11 @@ import {
     getResponseElements,
     getResponseContent,
     setupIntegrationTest,
-    dbOperationParameters,
+    databaseCollectionParameters,
     validateToolMetadata,
     validateAutoConnectBehavior,
     validateThrowsForInvalidArguments,
-    dbOperationInvalidArgTests,
+    databaseCollectionInvalidArgs,
 } from "../../../helpers.js";
 import { Document } from "bson";
 import { OptionalId } from "mongodb";
@@ -19,10 +19,10 @@ describe("collectionSchema tool", () => {
         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
@@ -1,10 +1,10 @@
 import {
     getResponseContent,
     setupIntegrationTest,
-    dbOperationParameters,
+    databaseCollectionParameters,
     validateToolMetadata,
     validateAutoConnectBehavior,
-    dbOperationInvalidArgTests,
+    databaseCollectionInvalidArgs,
     validateThrowsForInvalidArguments,
 } from "../../../helpers.js";
 import * as crypto from "crypto";
@@ -16,13 +16,13 @@ describe("collectionStorageSize tool", () => {
         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,99 @@
+import { ObjectId } from "bson";
+import {
+    setupIntegrationTest,
+    databaseParameters,
+    validateToolMetadata,
+    validateAutoConnectBehavior,
+    validateThrowsForInvalidArguments,
+    databaseInvalidArgs,
+    getResponseElements,
+} from "../../../helpers.js";
+import * as crypto from "crypto";
+
+describe("dbStats tool", () => {
+    const integration = setupIntegrationTest();
+
+    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,17 +5,21 @@ import {
     validateToolMetadata,
     validateAutoConnectBehavior,
     validateThrowsForInvalidArguments,
-    dbOperationInvalidArgTests,
+    databaseInvalidArgs,
+    databaseParameters,
 } from "../../../helpers.js";
 
 describe("listCollections tool", () => {
     const integration = setupIntegrationTest();
 
-    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
@@ -4,7 +4,6 @@ import {
     setupIntegrationTest,
     validateAutoConnectBehavior,
 } from "../../../helpers.js";
-import { toIncludeSameMembers } from "jest-extended";
 
 describe("listDatabases tool", () => {
     const integration = setupIntegrationTest();
@@ -26,7 +25,7 @@ describe("listDatabases tool", () => {
             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
