chore: refactor helpers to use vi.waitFor

himanshusinghs · himanshusinghs · commit 3e578580c711 · 2025-10-16T14:31:20.000+02:00
diff --git a/tests/integration/helpers.ts b/tests/integration/helpers.ts
@@ -1,4 +1,3 @@
-import type { Collection } from "mongodb";
 import { CompositeLogger } from "../../src/common/logger.js";
 import { ExportsManager } from "../../src/common/exportsManager.js";
 import { Session } from "../../src/common/session.js";
@@ -24,9 +23,6 @@ import { Elicitation } from "../../src/elicitation.js";
 import type { MockClientCapabilities, createMockElicitInput } from "../utils/elicitationMocks.js";
 import { VectorSearchEmbeddingsManager } from "../../src/common/search/vectorSearchEmbeddingsManager.js";
 
-export const DEFAULT_WAIT_TIMEOUT = 1000;
-export const DEFAULT_RETRY_INTERVAL = 100;
-
 export const driverOptions = setupDriverConfig({
     config,
     defaults: defaultDriverOptionsFromConfig,
@@ -423,60 +419,3 @@ export function getDataFromUntrustedContent(content: string): string {
 export function sleep(ms: number): Promise<void> {
     return new Promise((resolve) => setTimeout(resolve, ms));
 }
-
-export async function waitUntilSearchManagementServiceIsReady(
-    collection: Collection,
-    timeout: number = DEFAULT_WAIT_TIMEOUT,
-    interval: number = DEFAULT_RETRY_INTERVAL
-): Promise<void> {
-    await vi.waitFor(async () => await collection.listSearchIndexes({}).toArray(), { timeout, interval });
-}
-
-async function waitUntilSearchIndexIs(
-    collection: Collection,
-    searchIndex: string,
-    indexValidator: (index: { name: string; queryable: boolean }) => boolean,
-    timeout: number,
-    interval: number
-): Promise<void> {
-    await vi.waitFor(
-        async () => {
-            const searchIndexes = (await collection.listSearchIndexes(searchIndex).toArray()) as {
-                name: string;
-                queryable: boolean;
-            }[];
-
-            if (!searchIndexes.some((index) => indexValidator(index))) {
-                throw new Error("Search index did not pass validation");
-            }
-        },
-        {
-            timeout,
-            interval,
-        }
-    );
-}
-
-export async function waitUntilSearchIndexIsListed(
-    collection: Collection,
-    searchIndex: string,
-    timeout: number = DEFAULT_WAIT_TIMEOUT,
-    interval: number = DEFAULT_RETRY_INTERVAL
-): Promise<void> {
-    return waitUntilSearchIndexIs(collection, searchIndex, (index) => index.name === searchIndex, timeout, interval);
-}
-
-export async function waitUntilSearchIndexIsQueryable(
-    collection: Collection,
-    searchIndex: string,
-    timeout: number = DEFAULT_WAIT_TIMEOUT,
-    interval: number = DEFAULT_RETRY_INTERVAL
-): Promise<void> {
-    return waitUntilSearchIndexIs(
-        collection,
-        searchIndex,
-        (index) => index.name === searchIndex && index.queryable,
-        timeout,
-        interval
-    );
-}
diff --git a/tests/integration/tools/mongodb/create/createIndex.test.ts b/tests/integration/tools/mongodb/create/createIndex.test.ts
@@ -8,9 +8,8 @@ import {
     expectDefined,
     defaultTestConfig,
 } from "../../../helpers.js";
-import { ObjectId, type IndexDirection } from "mongodb";
-import { beforeEach, describe, expect, it } from "vitest";
-import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
+import { ObjectId, type Collection, type Document, type IndexDirection } from "mongodb";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
 
 describeWithMongoDB("createIndex tool when search is not enabled", (integration) => {
     it("doesn't allow creating vector search indexes", async () => {
@@ -403,12 +402,9 @@ describeWithMongoDB(
 describeWithMongoDB(
     "createIndex tool with vector search indexes",
     (integration) => {
-        let provider: NodeDriverServiceProvider;
-
-        beforeEach(async ({ signal }) => {
+        beforeEach(async () => {
             await integration.connectMcpClient();
-            provider = integration.mcpServer().session.serviceProvider;
-            await waitUntilSearchIsReady(provider, signal);
+            await waitUntilSearchIsReady(integration.mongoClient());
         });
 
         describe("when the collection does not exist", () => {
@@ -457,14 +453,26 @@ describeWithMongoDB(
         });
 
         describe("when the collection exists", () => {
+            let collectionName: string;
+            let collection: Collection;
+            beforeEach(async () => {
+                collectionName = new ObjectId().toString();
+                collection = await integration
+                    .mongoClient()
+                    .db(integration.randomDbName())
+                    .createCollection(collectionName);
+            });
+
+            afterEach(async () => {
+                await collection.drop();
+            });
+
             it("creates the index", async () => {
-                const collection = new ObjectId().toString();
-                await provider.createCollection(integration.randomDbName(), collection);
                 const response = await integration.mcpClient().callTool({
                     name: "create-index",
                     arguments: {
                         database: integration.randomDbName(),
-                        collection,
+                        collection: collectionName,
                         name: "vector_1_vector",
                         definition: [
                             {
@@ -480,10 +488,10 @@ describeWithMongoDB(
 
                 const content = getResponseContent(response.content);
                 expect(content).toEqual(
-                    `Created the index "vector_1_vector" on collection "${collection}" in database "${integration.randomDbName()}". Since this is a vector search index, it may take a while for the index to build. Use the \`list-indexes\` tool to check the index status.`
+                    `Created the index "vector_1_vector" on collection "${collectionName}" in database "${integration.randomDbName()}". Since this is a vector search index, it may take a while for the index to build. Use the \`list-indexes\` tool to check the index status.`
                 );
 
-                const indexes = await provider.getSearchIndexes(integration.randomDbName(), collection);
+                const indexes = (await collection.listSearchIndexes().toArray()) as unknown as Document[];
                 expect(indexes).toHaveLength(1);
                 expect(indexes[0]?.name).toEqual("vector_1_vector");
                 expect(indexes[0]?.type).toEqual("vectorSearch");
diff --git a/tests/integration/tools/mongodb/create/insertMany.test.ts b/tests/integration/tools/mongodb/create/insertMany.test.ts
@@ -1,7 +1,7 @@
 import {
-    createVectorSearchIndexAndWait,
     describeWithMongoDB,
     validateAutoConnectBehavior,
+    createVectorSearchIndexAndWait,
     waitUntilSearchIsReady,
 } from "../mongodbHelpers.js";
 
@@ -14,8 +14,8 @@ import {
     getDataFromUntrustedContent,
 } from "../../../helpers.js";
 import { beforeEach, afterEach, expect, it } from "vitest";
-import type { NodeDriverServiceProvider } from "@mongosh/service-provider-node-driver";
 import { ObjectId } from "bson";
+import type { Collection } from "mongodb";
 
 describeWithMongoDB("insertMany tool when search is disabled", (integration) => {
     validateToolMetadata(integration, "insert-many", "Insert an array of documents into a MongoDB collection", [
@@ -109,35 +109,29 @@ describeWithMongoDB("insertMany tool when search is disabled", (integration) =>
 describeWithMongoDB(
     "insertMany tool when search is enabled",
     (integration) => {
-        let provider: NodeDriverServiceProvider;
+        // let provider: NodeDriverServiceProvider;
+        let collection: Collection;
 
-        beforeEach(async ({ signal }) => {
+        beforeEach(async () => {
             await integration.connectMcpClient();
-            provider = integration.mcpServer().session.serviceProvider;
-            await provider.createCollection(integration.randomDbName(), "test");
-            await waitUntilSearchIsReady(provider, signal);
+            collection = await integration.mongoClient().db(integration.randomDbName()).createCollection("test");
+            await waitUntilSearchIsReady(integration.mongoClient());
         });
 
         afterEach(async () => {
-            await provider.dropCollection(integration.randomDbName(), "test");
+            await collection.drop();
         });
 
-        it("inserts a document when the embedding is correct", async ({ signal }) => {
-            await createVectorSearchIndexAndWait(
-                provider,
-                integration.randomDbName(),
-                "test",
-                [
-                    {
-                        type: "vector",
-                        path: "embedding",
-                        numDimensions: 8,
-                        similarity: "euclidean",
-                        quantization: "scalar",
-                    },
-                ],
-                signal
-            );
+        it("inserts a document when the embedding is correct", async () => {
+            await createVectorSearchIndexAndWait(integration.mongoClient(), integration.randomDbName(), "test", [
+                {
+                    type: "vector",
+                    path: "embedding",
+                    numDimensions: 8,
+                    similarity: "euclidean",
+                    quantization: "scalar",
+                },
+            ]);
 
             const response = await integration.mcpClient().callTool({
                 name: "insert-many",
@@ -152,26 +146,20 @@ describeWithMongoDB(
             const insertedIds = extractInsertedIds(content);
             expect(insertedIds).toHaveLength(1);
 
-            const docCount = await provider.countDocuments(integration.randomDbName(), "test", { _id: insertedIds[0] });
+            const docCount = await collection.countDocuments({ _id: insertedIds[0] });
             expect(docCount).toBe(1);
         });
 
-        it("returns an error when there is a search index and quantisation is wrong", async ({ signal }) => {
-            await createVectorSearchIndexAndWait(
-                provider,
-                integration.randomDbName(),
-                "test",
-                [
-                    {
-                        type: "vector",
-                        path: "embedding",
-                        numDimensions: 8,
-                        similarity: "euclidean",
-                        quantization: "scalar",
-                    },
-                ],
-                signal
-            );
+        it("returns an error when there is a search index and quantisation is wrong", async () => {
+            await createVectorSearchIndexAndWait(integration.mongoClient(), integration.randomDbName(), "test", [
+                {
+                    type: "vector",
+                    path: "embedding",
+                    numDimensions: 8,
+                    similarity: "euclidean",
+                    quantization: "scalar",
+                },
+            ]);
 
             const response = await integration.mcpClient().callTool({
                 name: "insert-many",
@@ -189,7 +177,7 @@ describeWithMongoDB(
                 "- Field embedding is an embedding with 8 dimensions and scalar quantization, and the provided value is not compatible."
             );
 
-            const oopsieCount = await provider.countDocuments(integration.randomDbName(), "test", {
+            const oopsieCount = await collection.countDocuments({
                 embedding: "oopsie",
             });
             expect(oopsieCount).toBe(0);
diff --git a/tests/integration/tools/mongodb/delete/dropIndex.test.ts b/tests/integration/tools/mongodb/delete/dropIndex.test.ts
@@ -8,15 +8,16 @@ import {
     getResponseContent,
     validateThrowsForInvalidArguments,
     validateToolMetadata,
-    waitUntilSearchIndexIsListed,
-    waitUntilSearchManagementServiceIsReady,
 } from "../../../helpers.js";
-import { describeWithMongoDB, type MongoDBIntegrationTestCase } from "../mongodbHelpers.js";
+import {
+    describeWithMongoDB,
+    waitUntilSearchIndexIsListed,
+    waitUntilSearchIsReady,
+    type MongoDBIntegrationTestCase,
+} from "../mongodbHelpers.js";
 import { createMockElicitInput } from "../../../../utils/elicitationMocks.js";
 import { Elicitation } from "../../../../../src/elicitation.js";
 
-const SEARCH_TIMEOUT = 20_000;
-
 function setupForClassicIndexes(integration: MongoDBIntegrationTestCase): {
     getMoviesCollection: () => Collection;
     getIndexName: () => string;
@@ -66,12 +67,12 @@ function setupForVectorSearchIndexes(integration: MongoDBIntegrationTestCase): {
                 plot: "This is a horrible movie about a database called BongoDB and how it tried to copy the OG MangoDB.",
             },
         ]);
-        await waitUntilSearchManagementServiceIsReady(moviesCollection, SEARCH_TIMEOUT);
+        await waitUntilSearchIsReady(mongoClient);
         await moviesCollection.createSearchIndex({
             name: indexName,
             definition: { mappings: { dynamic: true } },
         });
-        await waitUntilSearchIndexIsListed(moviesCollection, indexName, SEARCH_TIMEOUT);
+        await waitUntilSearchIndexIsListed(moviesCollection, indexName);
     });
 
     afterEach(async () => {
@@ -366,7 +367,7 @@ describe.each([{ vectorSearchEnabled: false }, { vectorSearchEnabled: true }])(
                         });
 
                         describe("and attempting to delete an existing index", () => {
-                            it("should succeed in deleting the index", { timeout: SEARCH_TIMEOUT }, async () => {
+                            it("should succeed in deleting the index", async () => {
                                 const response = await integration.mcpClient().callTool({
                                     name: "drop-index",
                                     arguments: {
diff --git a/tests/integration/tools/mongodb/mongodbHelpers.ts b/tests/integration/tools/mongodb/mongodbHelpers.ts
diff --git a/tests/integration/tools/mongodb/search/listSearchIndexes.test.ts b/tests/integration/tools/mongodb/search/listSearchIndexes.test.ts
