Skip to content

Commit 780ec67

Browse files
chore: use ListCollectionsTool response creators in tests
1 parent b9123ea commit 780ec67

File tree

2 files changed

+72
-35
lines changed

2 files changed

+72
-35
lines changed

src/tools/mongodb/metadata/listCollections.ts

Lines changed: 26 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,28 @@ import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
22
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
33
import { ToolArgs, OperationType } from "../../tool.js";
44

5+
export function listCollectionsResponse(database: string, collections: string[]): CallToolResult {
6+
if (collections.length === 0) {
7+
return {
8+
content: [
9+
{
10+
type: "text",
11+
text: `No collections found for database "${database}". To create a collection, use the "create-collection" tool.`,
12+
},
13+
],
14+
};
15+
}
16+
17+
return {
18+
content: collections.map((collection) => {
19+
return {
20+
text: `Name: "${collection}"`,
21+
type: "text",
22+
};
23+
}),
24+
};
25+
}
26+
527
export class ListCollectionsTool extends MongoDBToolBase {
628
protected name = "list-collections";
729
protected description = "List all collections for a given database";
@@ -15,24 +37,9 @@ export class ListCollectionsTool extends MongoDBToolBase {
1537
const provider = await this.ensureConnected();
1638
const collections = await provider.listCollections(database);
1739

18-
if (collections.length === 0) {
19-
return {
20-
content: [
21-
{
22-
type: "text",
23-
text: `No collections found for database "${database}". To create a collection, use the "create-collection" tool.`,
24-
},
25-
],
26-
};
27-
}
28-
29-
return {
30-
content: collections.map((collection) => {
31-
return {
32-
text: `Name: "${collection.name}"`,
33-
type: "text",
34-
};
35-
}),
36-
};
40+
return listCollectionsResponse(
41+
database,
42+
collections.map((collection) => `${collection.name}`)
43+
);
3744
}
3845
}
Lines changed: 46 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,16 @@
11
import { describeAccuracyTests } from "./sdk/describe-accuracy-tests.js";
22
import { getAvailableModels } from "./sdk/models.js";
33
import { AccuracyTestConfig } from "./sdk/describe-accuracy-tests.js";
4+
import { listCollectionsResponse } from "../../src/tools/mongodb/metadata/listCollections.js";
5+
import { listDatabasesResponse } from "../../src/tools/mongodb/metadata/listDatabases.js";
46

5-
function describeListCollectionsAccuracyTests(prompt: string): AccuracyTestConfig {
7+
function callsListCollections(prompt: string): AccuracyTestConfig {
68
return {
79
injectConnectedAssumption: true,
810
prompt: prompt,
911
mockedTools: {
1012
"list-collections": function listCollections() {
11-
return {
12-
content: [
13-
{
14-
type: "text",
15-
text: "Name: coll1",
16-
},
17-
{
18-
type: "text",
19-
text: "Name: coll1",
20-
},
21-
],
22-
};
13+
return listCollectionsResponse("db1", ["coll1", "coll2"]);
2314
},
2415
},
2516
expectedToolCalls: [
@@ -31,8 +22,47 @@ function describeListCollectionsAccuracyTests(prompt: string): AccuracyTestConfi
3122
};
3223
}
3324

25+
function callsListDatabasesAndListCollections(prompt: string): AccuracyTestConfig {
26+
return {
27+
injectConnectedAssumption: true,
28+
prompt: prompt,
29+
mockedTools: {
30+
"list-collections": function listCollections() {
31+
return listCollectionsResponse("db1", ["coll1", "coll2"]);
32+
},
33+
"list-databases": function listDatabases() {
34+
return listDatabasesResponse([
35+
{
36+
name: "db1",
37+
sizeOnDisk: "1024",
38+
},
39+
{
40+
name: "db2",
41+
sizeOnDisk: "2048",
42+
},
43+
]);
44+
},
45+
},
46+
expectedToolCalls: [
47+
{
48+
toolName: "list-databases",
49+
parameters: {},
50+
},
51+
{
52+
toolName: "list-collections",
53+
parameters: { database: "db1" },
54+
},
55+
{
56+
toolName: "list-collections",
57+
parameters: { database: "db2" },
58+
},
59+
],
60+
};
61+
}
62+
3463
describeAccuracyTests("list-collections", getAvailableModels(), [
35-
describeListCollectionsAccuracyTests("How many collections do I have in database db1?"),
36-
describeListCollectionsAccuracyTests("List all the collections in my MongoDB database db1."),
37-
describeListCollectionsAccuracyTests("Is there a coll1 collection in my MongoDB database db1?"),
64+
callsListCollections("How many collections do I have in database db1?"),
65+
callsListCollections("List all the collections in my MongoDB database db1."),
66+
callsListCollections("Is there a coll1 collection in my MongoDB database db1?"),
67+
callsListDatabasesAndListCollections("List all the collections that I have in total on my cluster?"),
3868
]);

0 commit comments

Comments
 (0)