Skip to content

Commit f1be251

Browse files
himanshusinghshimanshusinghs
committed
chore: tests for the new behavior
1 parent 250299b commit f1be251

File tree

4 files changed

+253
-28
lines changed

4 files changed

+253
-28
lines changed

‎src/tools/mongodb/read/aggregate.ts‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,7 @@ import { type Document, EJSON } from "bson";
1010
import { ErrorCodes, MongoDBError } from "../../../common/errors.js";
1111
import { iterateCursorUntilMaxBytes } from "../../../helpers/iterateCursor.js";
1212
import { operationWithFallback } from "../../../helpers/operationWithFallback.js";
13-
14-
/**
15-
* A cap for the maxTimeMS used for counting resulting documents of an
16-
* aggregation.
17-
*/
18-
const AGG_COUNT_MAX_TIME_MS_CAP = 60_000;
13+
import { AGG_COUNT_MAX_TIME_MS_CAP } from "../../../helpers/constants.js";
1914

2015
export const AggregateArgs = {
2116
pipeline: z.array(z.object({}).passthrough()).describe("An array of aggregation stages to execute"),

‎src/tools/mongodb/read/find.ts‎

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -8,15 +8,7 @@ import { checkIndexUsage } from "../../../helpers/indexCheck.js";
88
import { EJSON } from "bson";
99
import { iterateCursorUntilMaxBytes } from "../../../helpers/iterateCursor.js";
1010
import { operationWithFallback } from "../../../helpers/operationWithFallback.js";
11-
12-
/**
13-
* A cap for the maxTimeMS used for FindCursor.countDocuments.
14-
*
15-
* The number is relatively smaller because we expect the count documents query
16-
* to be finished sooner if not by the time the batch of documents is retrieved
17-
* so that count documents query don't hold the final response back.
18-
*/
19-
const QUERY_COUNT_MAX_TIME_MS_CAP = 10_000;
11+
import { QUERY_COUNT_MAX_TIME_MS_CAP } from "../../../helpers/constants.js";
2012

2113
export const FindArgs = {
2214
filter: z

‎tests/integration/tools/mongodb/read/aggregate.test.ts‎

Lines changed: 125 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,12 @@ import {
33
validateToolMetadata,
44
validateThrowsForInvalidArguments,
55
getResponseContent,
6+
defaultTestConfig,
67
} from "../../../helpers.js";
7-
import { expect, it } from "vitest";
8+
import { describe, expect, it, vi } from "vitest";
89
import { describeWithMongoDB, getDocsFromUntrustedContent, validateAutoConnectBehavior } from "../mongodbHelpers.js";
10+
import * as constants from "../../../../../src/helpers/constants.js";
11+
import { freshInsertDocuments } from "./find.test.js";
912

1013
describeWithMongoDB("aggregate tool", (integration) => {
1114
validateToolMetadata(integration, "aggregate", "Run an aggregation against a MongoDB collection", [
@@ -27,7 +30,7 @@ describeWithMongoDB("aggregate tool", (integration) => {
2730
{ database: 123, collection: "foo", pipeline: [] },
2831
]);
2932

30-
it("can run aggragation on non-existent database", async () => {
33+
it("can run aggregation on non-existent database", async () => {
3134
await integration.connectMcpClient();
3235
const response = await integration.mcpClient().callTool({
3336
name: "aggregate",
@@ -38,7 +41,7 @@ describeWithMongoDB("aggregate tool", (integration) => {
3841
expect(content).toEqual("The aggregation resulted in 0 documents.");
3942
});
4043

41-
it("can run aggragation on an empty collection", async () => {
44+
it("can run aggregation on an empty collection", async () => {
4245
await integration.mongoClient().db(integration.randomDbName()).createCollection("people");
4346

4447
await integration.connectMcpClient();
@@ -55,7 +58,7 @@ describeWithMongoDB("aggregate tool", (integration) => {
5558
expect(content).toEqual("The aggregation resulted in 0 documents.");
5659
});
5760

58-
it("can run aggragation on an existing collection", async () => {
61+
it("can run aggregation on an existing collection", async () => {
5962
const mongoClient = integration.mongoClient();
6063
await mongoClient
6164
.db(integration.randomDbName())
@@ -140,3 +143,121 @@ describeWithMongoDB("aggregate tool", (integration) => {
140143
};
141144
});
142145
});
146+
147+
describeWithMongoDB(
148+
"aggregate tool with configured max documents per query",
149+
(integration) => {
150+
describe("when the aggregation results are larger than the configured limit", () => {
151+
it("should return documents limited to the configured limit", async () => {
152+
await freshInsertDocuments({
153+
collection: integration.mongoClient().db(integration.randomDbName()).collection("people"),
154+
count: 1000,
155+
documentMapper(index) {
156+
return { name: `Person ${index}`, age: index };
157+
},
158+
});
159+
await integration.connectMcpClient();
160+
const response = await integration.mcpClient().callTool({
161+
name: "aggregate",
162+
arguments: {
163+
database: integration.randomDbName(),
164+
collection: "people",
165+
pipeline: [{ $match: { age: { $gte: 10 } } }, { $sort: { name: -1 } }],
166+
},
167+
});
168+
169+
const content = getResponseContent(response);
170+
expect(content).toContain("The aggregation resulted in 990 documents");
171+
expect(content).toContain(`Returning 20 documents while respecting the applied limits.`);
172+
const docs = getDocsFromUntrustedContent(content);
173+
expect(docs[0]).toEqual(
174+
expect.objectContaining({
175+
_id: expect.any(Object) as object,
176+
name: "Person 999",
177+
age: 999,
178+
})
179+
);
180+
expect(docs[1]).toEqual(
181+
expect.objectContaining({
182+
_id: expect.any(Object) as object,
183+
name: "Person 998",
184+
age: 998,
185+
})
186+
);
187+
});
188+
});
189+
190+
describe("when counting documents exceed the configured count maxTimeMS", () => {
191+
it("should abort discard count operation and respond with indeterminable count", async () => {
192+
vi.spyOn(constants, "AGG_COUNT_MAX_TIME_MS_CAP", "get").mockReturnValue(0.1);
193+
await freshInsertDocuments({
194+
collection: integration.mongoClient().db(integration.randomDbName()).collection("people"),
195+
count: 1000,
196+
documentMapper(index) {
197+
return { name: `Person ${index}`, age: index };
198+
},
199+
});
200+
await integration.connectMcpClient();
201+
const response = await integration.mcpClient().callTool({
202+
name: "aggregate",
203+
arguments: {
204+
database: integration.randomDbName(),
205+
collection: "people",
206+
pipeline: [{ $match: { age: { $gte: 10 } } }, { $sort: { name: -1 } }],
207+
},
208+
});
209+
const content = getResponseContent(response);
210+
expect(content).toContain("The aggregation resulted in indeterminable number of documents");
211+
expect(content).toContain(`Returning 20 documents while respecting the applied limits.`);
212+
const docs = getDocsFromUntrustedContent(content);
213+
expect(docs[0]).toEqual(
214+
expect.objectContaining({
215+
_id: expect.any(Object) as object,
216+
name: "Person 999",
217+
age: 999,
218+
})
219+
);
220+
expect(docs[1]).toEqual(
221+
expect.objectContaining({
222+
_id: expect.any(Object) as object,
223+
name: "Person 998",
224+
age: 998,
225+
})
226+
);
227+
vi.resetAllMocks();
228+
});
229+
});
230+
},
231+
() => ({ ...defaultTestConfig, maxDocumentsPerQuery: 20 })
232+
);
233+
234+
describeWithMongoDB(
235+
"aggregate tool with configured max bytes per query",
236+
(integration) => {
237+
describe("when the provided maxBytesPerQuery is hit", () => {
238+
it("should return only the documents that could fit in maxBytesPerQuery limit", async () => {
239+
await freshInsertDocuments({
240+
collection: integration.mongoClient().db(integration.randomDbName()).collection("people"),
241+
count: 1000,
242+
documentMapper(index) {
243+
return { name: `Person ${index}`, age: index };
244+
},
245+
});
246+
await integration.connectMcpClient();
247+
const response = await integration.mcpClient().callTool({
248+
name: "aggregate",
249+
arguments: {
250+
database: integration.randomDbName(),
251+
collection: "people",
252+
pipeline: [{ $match: { age: { $gte: 10 } } }, { $sort: { name: -1 } }],
253+
},
254+
});
255+
256+
const content = getResponseContent(response);
257+
expect(content).toContain("The aggregation resulted in 990 documents");
258+
expect(content).toContain(`Returning 1 documents while respecting the applied limits.`);
259+
});
260+
});
261+
},
262+
() => ({ ...defaultTestConfig, maxBytesPerQuery: 100 })
263+
);

‎tests/integration/tools/mongodb/read/find.test.ts‎

Lines changed: 126 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,30 @@
1-
import { beforeEach, describe, expect, it } from "vitest";
1+
import { beforeEach, describe, expect, it, vi } from "vitest";
2+
import type { Document, Collection } from "mongodb";
23
import {
34
getResponseContent,
45
databaseCollectionParameters,
56
validateToolMetadata,
67
validateThrowsForInvalidArguments,
78
expectDefined,
9+
defaultTestConfig,
810
} from "../../../helpers.js";
11+
import * as constants from "../../../../../src/helpers/constants.js";
912
import { describeWithMongoDB, getDocsFromUntrustedContent, validateAutoConnectBehavior } from "../mongodbHelpers.js";
1013

14+
export async function freshInsertDocuments({
15+
collection,
16+
count,
17+
documentMapper = (index): Document => ({ value: index }),
18+
}: {
19+
collection: Collection<Document>;
20+
count: number;
21+
documentMapper?: (index: number) => Document;
22+
}): Promise<void> {
23+
await collection.drop();
24+
const documents = Array.from({ length: count }).map((_, idx) => documentMapper(idx));
25+
await collection.insertMany(documents);
26+
}
27+
1128
describeWithMongoDB("find tool", (integration) => {
1229
validateToolMetadata(integration, "find", "Run a find query against a MongoDB collection", [
1330
...databaseCollectionParameters,
@@ -73,14 +90,10 @@ describeWithMongoDB("find tool", (integration) => {
7390

7491
describe("with existing database", () => {
7592
beforeEach(async () => {
76-
const mongoClient = integration.mongoClient();
77-
const items = Array(10)
78-
.fill(0)
79-
.map((_, index) => ({
80-
value: index,
81-
}));
82-
83-
await mongoClient.db(integration.randomDbName()).collection("foo").insertMany(items);
93+
await freshInsertDocuments({
94+
collection: integration.mongoClient().db(integration.randomDbName()).collection("foo"),
95+
count: 10,
96+
});
8497
});
8598

8699
const testCases: {
@@ -211,3 +224,107 @@ describeWithMongoDB("find tool", (integration) => {
211224
};
212225
});
213226
});
227+
228+
describeWithMongoDB(
229+
"find tool with configured max documents per query",
230+
(integration) => {
231+
describe("when the provided limit is lower than the configured max limit", () => {
232+
it("should return documents limited to the provided limit", async () => {
233+
await freshInsertDocuments({
234+
collection: integration.mongoClient().db(integration.randomDbName()).collection("foo"),
235+
count: 1000,
236+
});
237+
await integration.connectMcpClient();
238+
const response = await integration.mcpClient().callTool({
239+
name: "find",
240+
arguments: {
241+
database: integration.randomDbName(),
242+
collection: "foo",
243+
filter: {},
244+
// default is 10
245+
limit: undefined,
246+
},
247+
});
248+
249+
const content = getResponseContent(response);
250+
expect(content).toContain(`Query on collection "foo" resulted in 10 documents.`);
251+
expect(content).toContain(`Returning 10 documents while respecting the applied limits.`);
252+
});
253+
});
254+
255+
describe("when the provided limit is larger than the configured max limit", () => {
256+
it("should return documents limited to the configured max limit", async () => {
257+
await freshInsertDocuments({
258+
collection: integration.mongoClient().db(integration.randomDbName()).collection("foo"),
259+
count: 1000,
260+
});
261+
262+
await integration.connectMcpClient();
263+
const response = await integration.mcpClient().callTool({
264+
name: "find",
265+
arguments: {
266+
database: integration.randomDbName(),
267+
collection: "foo",
268+
filter: {},
269+
limit: 10000,
270+
},
271+
});
272+
273+
const content = getResponseContent(response);
274+
expect(content).toContain(`Query on collection "foo" resulted in 1000 documents.`);
275+
expect(content).toContain(`Returning 20 documents while respecting the applied limits.`);
276+
});
277+
});
278+
279+
describe("when counting documents exceed the configured count maxTimeMS", () => {
280+
it("should abort discard count operation and respond with indeterminable count", async () => {
281+
vi.spyOn(constants, "QUERY_COUNT_MAX_TIME_MS_CAP", "get").mockReturnValue(0.1);
282+
await freshInsertDocuments({
283+
collection: integration.mongoClient().db(integration.randomDbName()).collection("foo"),
284+
count: 1000,
285+
});
286+
await integration.connectMcpClient();
287+
const response = await integration.mcpClient().callTool({
288+
name: "find",
289+
arguments: { database: integration.randomDbName(), collection: "foo" },
290+
});
291+
const content = getResponseContent(response);
292+
expect(content).toContain('Query on collection "foo" resulted in indeterminable number of documents.');
293+
294+
const docs = getDocsFromUntrustedContent(content);
295+
expect(docs.length).toEqual(10);
296+
vi.resetAllMocks();
297+
});
298+
});
299+
},
300+
() => ({ ...defaultTestConfig, maxDocumentsPerQuery: 20 })
301+
);
302+
303+
describeWithMongoDB(
304+
"find tool with configured max bytes per query",
305+
(integration) => {
306+
describe("when the provided maxBytesPerQuery is hit", () => {
307+
it("should return only the documents that could fit in maxBytesPerQuery limit", async () => {
308+
await freshInsertDocuments({
309+
collection: integration.mongoClient().db(integration.randomDbName()).collection("foo"),
310+
count: 1000,
311+
});
312+
await integration.connectMcpClient();
313+
const response = await integration.mcpClient().callTool({
314+
name: "find",
315+
arguments: {
316+
database: integration.randomDbName(),
317+
collection: "foo",
318+
filter: {},
319+
limit: 1000,
320+
},
321+
});
322+
323+
const content = getResponseContent(response);
324+
expect(content).toContain(`Query on collection "foo" resulted in 1000 documents.`);
325+
expect(content).toContain(`Returning 1 documents while respecting the applied limits.`);
326+
});
327+
});
328+
},
329+
() => ({ ...defaultTestConfig, maxBytesPerQuery: 50 })
330+
);

0 commit comments

Comments
 (0)