chore: fix existing tests

Author: himanshusinghs

src/tools/mongodb/read/aggregate.ts

@@ -60,14 +60,21 @@ export class AggregateTool extends MongoDBToolBase {
                 iterateCursorUntilMaxBytes(aggregationCursor, this.config.maxBytesPerQuery),
             ]);
 
-            const messageDescription = `\
-    The aggregation resulted in ${totalDocuments === undefined ? "indeterminable number of" : totalDocuments} documents. \
-    Returning ${documents.length} documents while respecting the applied limits. \
-    Note to LLM: If entire aggregation result is needed then use "export" tool to export the aggregation results.\
-    `;
+            let messageDescription = `\
+The aggregation resulted in ${totalDocuments === undefined ? "indeterminable number of" : totalDocuments} documents.\
+`;
+            if (documents.length) {
+                messageDescription += ` \
+Returning ${documents.length} documents while respecting the applied limits. \
+Note to LLM: If entire aggregation result is needed then use "export" tool to export the aggregation results.\
+`;
+            }
 
             return {
-                content: formatUntrustedData(messageDescription, EJSON.stringify(documents)),
+                content: formatUntrustedData(
+                    messageDescription,
+                    documents.length > 0 ? EJSON.stringify(documents) : undefined
+                ),
             };
         } finally {
             await aggregationCursor?.close();
@@ -114,7 +121,7 @@ export class AggregateTool extends MongoDBToolBase {
                 "totalDocuments" in documentWithCount &&
                 typeof documentWithCount.totalDocuments === "number"
                     ? documentWithCount.totalDocuments
-                    : undefined;
+                    : 0;
 
             return totalDocuments;
         }, undefined);

src/tools/mongodb/read/find.ts

@@ -89,14 +89,21 @@ export class FindTool extends MongoDBToolBase {
                 iterateCursorUntilMaxBytes(findCursor, this.config.maxBytesPerQuery),
             ]);
 
-            const messageDescription = `\
-Query on collection "${collection}" resulted in ${queryResultsCount === undefined ? "indeterminable number of" : queryResultsCount} documents. \
+            let messageDescription = `\
+Query on collection "${collection}" resulted in ${queryResultsCount === undefined ? "indeterminable number of" : queryResultsCount} documents.\
+`;
+            if (documents.length) {
+                messageDescription += ` \
 Returning ${documents.length} documents while respecting the applied limits. \
 Note to LLM: If entire query result is needed then use "export" tool to export the query results.\
 `;
+            }
 
             return {
-                content: formatUntrustedData(messageDescription, EJSON.stringify(documents)),
+                content: formatUntrustedData(
+                    messageDescription,
+                    documents.length > 0 ? EJSON.stringify(documents) : undefined
+                ),
             };
         } finally {
             await findCursor?.close();

tests/integration/indexCheck.test.ts

@@ -61,8 +61,7 @@ describe("IndexCheck integration tests", () => {
 
                         expect(response.isError).toBeFalsy();
                         const content = getResponseContent(response.content);
-                        expect(content).toContain("Found");
-                        expect(content).toContain("documents");
+                        expect(content).toContain('Query on collection "find-test-collection" resulted in');
                     });
 
                     it("should allow queries using _id (IDHACK)", async () => {
@@ -86,7 +85,9 @@ describe("IndexCheck integration tests", () => {
 
                         expect(response.isError).toBeFalsy();
                         const content = getResponseContent(response.content);
-                        expect(content).toContain("Found 1 documents");
+                        expect(content).toContain(
+                            'Query on collection "find-test-collection" resulted in 1 documents.'
+                        );
                     });
                 });
 
@@ -351,7 +352,7 @@ describe("IndexCheck integration tests", () => {
 
                     expect(findResponse.isError).toBeFalsy();
                     const findContent = getResponseContent(findResponse.content);
-                    expect(findContent).toContain("Found");
+                    expect(findContent).toContain('Query on collection "disabled-test-collection" resulted in');
                     expect(findContent).not.toContain("Index check failed");
                 });
 

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

@@ -56,7 +56,7 @@ describeWithMongoDB("find tool", (integration) => {
             arguments: { database: "non-existent", collection: "foos" },
         });
         const content = getResponseContent(response.content);
-        expect(content).toEqual('Found 0 documents in the collection "foos".');
+        expect(content).toEqual('Query on collection "foos" resulted in 0 documents.');
     });
 
     it("returns 0 when collection doesn't exist", async () => {
@@ -68,7 +68,7 @@ describeWithMongoDB("find tool", (integration) => {
             arguments: { database: integration.randomDbName(), collection: "non-existent" },
         });
         const content = getResponseContent(response.content);
-        expect(content).toEqual('Found 0 documents in the collection "non-existent".');
+        expect(content).toEqual('Query on collection "non-existent" resulted in 0 documents.');
     });
 
     describe("with existing database", () => {
@@ -148,7 +148,7 @@ describeWithMongoDB("find tool", (integration) => {
                     },
                 });
                 const content = getResponseContent(response);
-                expect(content).toContain(`Found ${expected.length} documents in the collection "foo".`);
+                expect(content).toContain(`Query on collection "foo" resulted in ${expected.length} documents.`);
 
                 const docs = getDocsFromUntrustedContent(content);
 
@@ -165,7 +165,7 @@ describeWithMongoDB("find tool", (integration) => {
                 arguments: { database: integration.randomDbName(), collection: "foo" },
             });
             const content = getResponseContent(response);
-            expect(content).toContain('Found 10 documents in the collection "foo".');
+            expect(content).toContain('Query on collection "foo" resulted in 10 documents.');
 
             const docs = getDocsFromUntrustedContent(content);
             expect(docs.length).toEqual(10);
@@ -195,7 +195,7 @@ describeWithMongoDB("find tool", (integration) => {
             });
 
             const content = getResponseContent(response);
-            expect(content).toContain('Found 1 documents in the collection "foo".');
+            expect(content).toContain('Query on collection "foo" resulted in 1 documents.');
 
             const docs = getDocsFromUntrustedContent(content);
             expect(docs.length).toEqual(1);
@@ -207,7 +207,7 @@ describeWithMongoDB("find tool", (integration) => {
     validateAutoConnectBehavior(integration, "find", () => {
         return {
             args: { database: integration.randomDbName(), collection: "coll1" },
-            expectedResponse: 'Found 0 documents in the collection "coll1"',
+            expectedResponse: 'Query on collection "coll1" resulted in 0 documents.',
         };
     });
 });