diff --git a/src/tools/mongodb/metadata/dbStats.ts b/src/tools/mongodb/metadata/dbStats.ts index 979b17e2..a8c0ea0d 100644 --- 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 index 243f6a7b..5b7ebe1c 100644 --- 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 index c720dbe1..ef8da5f1 100644 --- a/tests/integration/tools/mongodb/create/createCollection.test.ts +++ b/tests/integration/tools/mongodb/create/createCollection.test.ts @@ -2,10 +2,10 @@ import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelp import { getResponseContent, - dbOperationParameters, + databaseCollectionParameters, validateToolMetadata, validateThrowsForInvalidArguments, - dbOperationInvalidArgTests, + databaseCollectionInvalidArgs, } from "../../../helpers.js"; describeWithMongoDB("createCollection tool", (integration) => { @@ -13,10 +13,10 @@ describeWithMongoDB("createCollection tool", (integration) => { 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 index 83687f3d..fa921339 100644 --- a/tests/integration/tools/mongodb/create/createIndex.test.ts +++ b/tests/integration/tools/mongodb/create/createIndex.test.ts @@ -2,7 +2,7 @@ import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelp import { getResponseContent, - dbOperationParameters, + databaseCollectionParameters, validateToolMetadata, validateThrowsForInvalidArguments, } from "../../../helpers.js"; @@ -10,7 +10,7 @@ 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 index e05b3563..b4042029 100644 --- 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 index 5f8e096b..9201d566 100644 --- 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 index 764f58e5..1dcaa218 100644 --- a/tests/integration/tools/mongodb/delete/dropCollection.test.ts +++ b/tests/integration/tools/mongodb/delete/dropCollection.test.ts @@ -2,10 +2,10 @@ import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelp import { getResponseContent, - dbOperationParameters, + databaseCollectionParameters, validateToolMetadata, validateThrowsForInvalidArguments, - dbOperationInvalidArgTests, + databaseCollectionInvalidArgs, } from "../../../helpers.js"; describeWithMongoDB("dropCollection tool", (integration) => { @@ -13,10 +13,10 @@ describeWithMongoDB("dropCollection tool", (integration) => { 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 index a7fb7d92..29a79206 100644 --- a/tests/integration/tools/mongodb/delete/dropDatabase.test.ts +++ b/tests/integration/tools/mongodb/delete/dropDatabase.test.ts @@ -2,10 +2,10 @@ import { describeWithMongoDB, validateAutoConnectBehavior } from "../mongodbHelp import { getResponseContent, - dbOperationParameters, validateToolMetadata, validateThrowsForInvalidArguments, - dbOperationInvalidArgTests, + databaseParameters, + databaseInvalidArgs, } from "../../../helpers.js"; describeWithMongoDB("dropDatabase tool", (integration) => { @@ -13,10 +13,10 @@ describeWithMongoDB("dropDatabase tool", (integration) => { 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 index b9d14160..ccfc988f 100644 --- 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 index bc7c35ca..23e86cde 100644 --- 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 new file mode 100644 index 00000000..8e4a57c7 --- /dev/null +++ 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 index 15904874..cef0a59d 100644 --- 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 index 27984249..3288cf30 100644 --- 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 index 48d31077..938285a8 100644 --- 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", [