Skip to content

Commit 5d378cc

Browse files
nirinchevnirinchev
authored
chore: add tests for metadata actions (#91)
1 parent 9ba243d commit 5d378cc

17 files changed

+638
-578
lines changed

‎src/tools/mongodb/metadata/collectionSchema.ts

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
22
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
33
import { ToolArgs, OperationType } from "../../tool.js";
4-
import { parseSchema, SchemaField } from "mongodb-schema";
4+
import { getSimplifiedSchema } from "mongodb-schema";
55

66
export class CollectionSchemaTool extends MongoDBToolBase {
77
protected name = "collection-schema";
@@ -13,29 +13,31 @@ export class CollectionSchemaTool extends MongoDBToolBase {
1313
protected async execute({ database, collection }: ToolArgs<typeof DbOperationArgs>): Promise<CallToolResult> {
1414
const provider = await this.ensureConnected();
1515
const documents = await provider.find(database, collection, {}, { limit: 5 }).toArray();
16-
const schema = await parseSchema(documents);
16+
const schema = await getSimplifiedSchema(documents);
17+
18+
const fieldsCount = Object.entries(schema).length;
19+
if (fieldsCount === 0) {
20+
return {
21+
content: [
22+
{
23+
text: `Could not deduce the schema for "${database}.${collection}". This may be because it doesn't exist or is empty.`,
24+
type: "text",
25+
},
26+
],
27+
};
28+
}
1729

1830
return {
1931
content: [
2032
{
21-
text: `Found ${schema.fields.length} fields in the schema for \`${database}.${collection}\``,
33+
text: `Found ${fieldsCount} fields in the schema for "${database}.${collection}"`,
2234
type: "text",
2335
},
2436
{
25-
text: this.formatFieldOutput(schema.fields),
37+
text: JSON.stringify(schema),
2638
type: "text",
2739
},
2840
],
2941
};
3042
}
31-
32-
private formatFieldOutput(fields: SchemaField[]): string {
33-
let result = "| Field | Type | Confidence |\n";
34-
result += "|-------|------|-------------|\n";
35-
for (const field of fields) {
36-
const fieldType = Array.isArray(field.type) ? field.type.join(", ") : field.type;
37-
result += `| ${field.name} | \`${fieldType}\` | ${(field.probability * 100).toFixed(0)}% |\n`;
38-
}
39-
return result;
40-
}
4143
}

‎src/tools/mongodb/metadata/collectionStorageSize.ts

Lines changed: 41 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import { ToolArgs, OperationType } from "../../tool.js";
44

55
export class CollectionStorageSizeTool extends MongoDBToolBase {
66
protected name = "collection-storage-size";
7-
protected description = "Gets the size of the collection in MB";
7+
protected description = "Gets the size of the collection";
88
protected argsShape = DbOperationArgs;
99

1010
protected operationType: OperationType = "metadata";
@@ -14,17 +14,55 @@ export class CollectionStorageSizeTool extends MongoDBToolBase {
1414
const [{ value }] = (await provider
1515
.aggregate(database, collection, [
1616
{ $collStats: { storageStats: {} } },
17-
{ $group: { _id: null, value: { $sum: "$storageStats.storageSize" } } },
17+
{ $group: { _id: null, value: { $sum: "$storageStats.size" } } },
1818
])
1919
.toArray()) as [{ value: number }];
2020

21+
const { units, value: scaledValue } = CollectionStorageSizeTool.getStats(value);
22+
2123
return {
2224
content: [
2325
{
24-
text: `The size of \`${database}.${collection}\` is \`${(value / 1024 / 1024).toFixed(2)} MB\``,
26+
text: `The size of "${database}.${collection}" is \`${scaledValue.toFixed(2)} ${units}\``,
2527
type: "text",
2628
},
2729
],
2830
};
2931
}
32+
33+
protected handleError(
34+
error: unknown,
35+
args: ToolArgs<typeof this.argsShape>
36+
): Promise<CallToolResult> | CallToolResult {
37+
if (error instanceof Error && "codeName" in error && error.codeName === "NamespaceNotFound") {
38+
return {
39+
content: [
40+
{
41+
text: `The size of "${args.database}.${args.collection}" cannot be determined because the collection does not exist.`,
42+
type: "text",
43+
},
44+
],
45+
};
46+
}
47+
48+
return super.handleError(error, args);
49+
}
50+
51+
private static getStats(value: number): { value: number; units: string } {
52+
const kb = 1024;
53+
const mb = kb * 1024;
54+
const gb = mb * 1024;
55+
56+
if (value > gb) {
57+
return { value: value / gb, units: "GB" };
58+
}
59+
60+
if (value > mb) {
61+
return { value: value / mb, units: "MB" };
62+
}
63+
if (value > kb) {
64+
return { value: value / kb, units: "KB" };
65+
}
66+
return { value, units: "bytes" };
67+
}
3068
}

‎src/tools/mongodb/mongodbTool.ts

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { z } from "zod";
2-
import { ToolBase, ToolCategory } from "../tool.js";
2+
import { ToolArgs, ToolBase, ToolCategory } from "../tool.js";
33
import { Session } from "../../session.js";
44
import { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
55
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
@@ -30,7 +30,10 @@ export abstract class MongoDBToolBase extends ToolBase {
3030
return this.session.serviceProvider;
3131
}
3232

33-
protected handleError(error: unknown): Promise<CallToolResult> | CallToolResult {
33+
protected handleError(
34+
error: unknown,
35+
args: ToolArgs<typeof this.argsShape>
36+
): Promise<CallToolResult> | CallToolResult {
3437
if (error instanceof MongoDBError && error.code === ErrorCodes.NotConnectedToMongoDB) {
3538
return {
3639
content: [
@@ -47,7 +50,7 @@ export abstract class MongoDBToolBase extends ToolBase {
4750
};
4851
}
4952

50-
return super.handleError(error);
53+
return super.handleError(error, args);
5154
}
5255

5356
protected async connectToMongoDB(connectionString: string): Promise<void> {

‎src/tools/tool.ts

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ export abstract class ToolBase {
4444
} catch (error: unknown) {
4545
logger.error(mongoLogId(1_000_000), "tool", `Error executing ${this.name}: ${error as string}`);
4646

47-
return await this.handleError(error);
47+
return await this.handleError(error, args[0] as ToolArgs<typeof this.argsShape>);
4848
}
4949
};
5050

@@ -76,7 +76,11 @@ export abstract class ToolBase {
7676
}
7777

7878
// 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 {
8084
return {
8185
content: [
8286
{

‎tests/integration/helpers.ts

Lines changed: 88 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
99
import { MongoClient, ObjectId } from "mongodb";
1010
import { toIncludeAllMembers } from "jest-extended";
1111
import config from "../../src/config.js";
12+
import { McpError } from "@modelcontextprotocol/sdk/types.js";
1213

1314
interface ParameterInfo {
1415
name: string;
@@ -226,10 +227,93 @@ export const dbOperationParameters: ParameterInfo[] = [
226227
{ name: "collection", type: "string", description: "Collection name", required: true },
227228
];
228229

229-
export function validateParameters(tool: ToolInfo, parameters: ParameterInfo[]): void {
230-
const toolParameters = getParameters(tool);
231-
expect(toolParameters).toHaveLength(parameters.length);
232-
expect(toolParameters).toIncludeAllMembers(parameters);
230+
export const dbOperationInvalidArgTests = [{}, { database: 123 }, { foo: "bar", database: "test" }, { database: [] }];
231+
232+
export function validateToolMetadata(
233+
integration: IntegrationTest,
234+
name: string,
235+
description: string,
236+
parameters: ParameterInfo[]
237+
): void {
238+
it("should have correct metadata", async () => {
239+
const { tools } = await integration.mcpClient().listTools();
240+
const tool = tools.find((tool) => tool.name === name)!;
241+
expect(tool).toBeDefined();
242+
expect(tool.description).toBe(description);
243+
244+
const toolParameters = getParameters(tool);
245+
expect(toolParameters).toHaveLength(parameters.length);
246+
expect(toolParameters).toIncludeAllMembers(parameters);
247+
});
248+
}
249+
250+
export function validateAutoConnectBehavior(
251+
integration: IntegrationTest,
252+
name: string,
253+
validation: () => {
254+
args: { [x: string]: unknown };
255+
expectedResponse?: string;
256+
validate?: (content: unknown) => void;
257+
},
258+
beforeEachImpl?: () => Promise<void>
259+
): void {
260+
describe("when not connected", () => {
261+
if (beforeEachImpl) {
262+
beforeEach(() => beforeEachImpl());
263+
}
264+
265+
it("connects automatically if connection string is configured", async () => {
266+
config.connectionString = integration.connectionString();
267+
268+
const validationInfo = validation();
269+
270+
const response = await integration.mcpClient().callTool({
271+
name,
272+
arguments: validationInfo.args,
273+
});
274+
275+
if (validationInfo.expectedResponse) {
276+
const content = getResponseContent(response.content);
277+
expect(content).toContain(validationInfo.expectedResponse);
278+
}
279+
280+
if (validationInfo.validate) {
281+
validationInfo.validate(response.content);
282+
}
283+
});
284+
285+
it("throws an error if connection string is not configured", async () => {
286+
const response = await integration.mcpClient().callTool({
287+
name,
288+
arguments: validation().args,
289+
});
290+
const content = getResponseContent(response.content);
291+
expect(content).toContain("You need to connect to a MongoDB instance before you can access its data.");
292+
});
293+
});
294+
}
295+
296+
export function validateThrowsForInvalidArguments(
297+
integration: IntegrationTest,
298+
name: string,
299+
args: { [x: string]: unknown }[]
300+
): void {
301+
describe("with invalid arguments", () => {
302+
for (const arg of args) {
303+
it(`throws a schema error for: ${JSON.stringify(arg)}`, async () => {
304+
await integration.connectMcpClient();
305+
try {
306+
await integration.mcpClient().callTool({ name, arguments: arg });
307+
expect.fail("Expected an error to be thrown");
308+
} catch (error) {
309+
expect(error).toBeInstanceOf(McpError);
310+
const mcpError = error as McpError;
311+
expect(mcpError.code).toEqual(-32602);
312+
expect(mcpError.message).toContain(`Invalid arguments for tool ${name}`);
313+
}
314+
});
315+
}
316+
});
233317
}
234318

235319
export function describeAtlas(name: number | string | Function | jest.FunctionLike, fn: jest.EmptyFunction) {

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

Lines changed: 16 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -1,50 +1,24 @@
11
import {
22
getResponseContent,
3-
validateParameters,
43
dbOperationParameters,
54
setupIntegrationTest,
5+
validateToolMetadata,
6+
validateAutoConnectBehavior,
7+
validateThrowsForInvalidArguments,
8+
dbOperationInvalidArgTests,
69
} from "../../../helpers.js";
7-
import { toIncludeSameMembers } from "jest-extended";
8-
import { McpError } from "@modelcontextprotocol/sdk/types.js";
9-
import { ObjectId } from "bson";
10-
import config from "../../../../../src/config.js";
1110

1211
describe("createCollection tool", () => {
1312
const integration = setupIntegrationTest();
1413

15-
it("should have correct metadata", async () => {
16-
const { tools } = await integration.mcpClient().listTools();
17-
const listCollections = tools.find((tool) => tool.name === "create-collection")!;
18-
expect(listCollections).toBeDefined();
19-
expect(listCollections.description).toBe(
20-
"Creates a new collection in a database. If the database doesn't exist, it will be created automatically."
21-
);
14+
validateToolMetadata(
15+
integration,
16+
"create-collection",
17+
"Creates a new collection in a database. If the database doesn't exist, it will be created automatically.",
18+
dbOperationParameters
19+
);
2220

23-
validateParameters(listCollections, dbOperationParameters);
24-
});
25-
26-
describe("with invalid arguments", () => {
27-
const args = [
28-
{},
29-
{ database: 123, collection: "bar" },
30-
{ foo: "bar", database: "test", collection: "bar" },
31-
{ collection: [], database: "test" },
32-
];
33-
for (const arg of args) {
34-
it(`throws a schema error for: ${JSON.stringify(arg)}`, async () => {
35-
await integration.connectMcpClient();
36-
try {
37-
await integration.mcpClient().callTool({ name: "create-collection", arguments: arg });
38-
expect.fail("Expected an error to be thrown");
39-
} catch (error) {
40-
expect(error).toBeInstanceOf(McpError);
41-
const mcpError = error as McpError;
42-
expect(mcpError.code).toEqual(-32602);
43-
expect(mcpError.message).toContain("Invalid arguments for tool create-collection");
44-
}
45-
});
46-
}
47-
});
21+
validateThrowsForInvalidArguments(integration, "create-collection", dbOperationInvalidArgTests);
4822

4923
describe("with non-existent database", () => {
5024
it("creates a new collection", async () => {
@@ -114,25 +88,10 @@ describe("createCollection tool", () => {
11488
});
11589
});
11690

117-
describe("when not connected", () => {
118-
it("connects automatically if connection string is configured", async () => {
119-
config.connectionString = integration.connectionString();
120-
121-
const response = await integration.mcpClient().callTool({
122-
name: "create-collection",
123-
arguments: { database: integration.randomDbName(), collection: "new-collection" },
124-
});
125-
const content = getResponseContent(response.content);
126-
expect(content).toEqual(`Collection "new-collection" created in database "${integration.randomDbName()}".`);
127-
});
128-
129-
it("throws an error if connection string is not configured", async () => {
130-
const response = await integration.mcpClient().callTool({
131-
name: "create-collection",
132-
arguments: { database: integration.randomDbName(), collection: "new-collection" },
133-
});
134-
const content = getResponseContent(response.content);
135-
expect(content).toContain("You need to connect to a MongoDB instance before you can access its data.");
136-
});
91+
validateAutoConnectBehavior(integration, "create-collection", () => {
92+
return {
93+
args: { database: integration.randomDbName(), collection: "new-collection" },
94+
expectedResponse: `Collection "new-collection" created in database "${integration.randomDbName()}".`,
95+
};
13796
});
13897
});

0 commit comments

Comments
 (0)