chore: replace with vi.waitFor

Author: himanshusinghs

tests/integration/helpers.ts

@@ -23,8 +23,8 @@ import { Keychain } from "../../src/common/keychain.js";
 import { Elicitation } from "../../src/elicitation.js";
 import type { MockClientCapabilities, createMockElicitInput } from "../utils/elicitationMocks.js";
 
-const SEARCH_READY_CHECK_RETRIES = 200;
-const SEARCH_INDEX_STATUS_CHECK_RETRIES = 100;
+export const DEFAULT_WAIT_TIMEOUT = 1000;
+export const DEFAULT_RETRY_INTERVAL = 100;
 
 export const driverOptions = setupDriverConfig({
     config,
@@ -422,109 +422,59 @@ export function sleep(ms: number): Promise<void> {
     return new Promise((resolve) => setTimeout(resolve, ms));
 }
 
-export async function waitFor(
-    condition: () => boolean | Promise<boolean>,
-    {
-        retries,
-        retryTimeout,
-        abortSignal,
-        shouldRetryOnError,
-    }: {
-        retries: number;
-        retryTimeout: number;
-        abortSignal?: AbortSignal;
-        shouldRetryOnError?: (error: unknown) => boolean | Promise<boolean>;
-    } = {
-        retries: 100,
-        retryTimeout: 100,
-    }
-): Promise<void> {
-    for (let i = 0; i < retries && !abortSignal?.aborted; i++) {
-        try {
-            if (await condition()) {
-                return;
-            }
-
-            await sleep(retryTimeout);
-        } catch (error) {
-            if (shouldRetryOnError && (await shouldRetryOnError(error))) {
-                await sleep(retryTimeout);
-                continue;
-            }
-            throw error;
-        }
-    }
-}
-
 export async function waitUntilSearchManagementServiceIsReady(
     collection: Collection,
-    abortSignal?: AbortSignal
+    timeout: number = DEFAULT_WAIT_TIMEOUT,
+    interval: number = DEFAULT_RETRY_INTERVAL
 ): Promise<void> {
-    await waitFor(
-        async (): Promise<boolean> => {
-            await collection.listSearchIndexes({}).toArray();
-            return true;
-        },
-        {
-            retries: SEARCH_READY_CHECK_RETRIES,
-            retryTimeout: 100,
-            abortSignal,
-            shouldRetryOnError: (error: unknown) => {
-                return (
-                    error instanceof Error &&
-                    error.message.includes("Error connecting to Search Index Management service")
-                );
-            },
-        }
-    );
+    await vi.waitFor(async () => await collection.listSearchIndexes({}).toArray(), { timeout, interval });
 }
 
 async function waitUntilSearchIndexIs(
     collection: Collection,
     searchIndex: string,
     indexValidator: (index: { name: string; queryable: boolean }) => boolean,
-    abortSignal?: AbortSignal
+    timeout: number,
+    interval: number
 ): Promise<void> {
-    await waitFor(
-        async (): Promise<boolean> => {
+    await vi.waitFor(
+        async () => {
             const searchIndexes = (await collection.listSearchIndexes(searchIndex).toArray()) as {
                 name: string;
                 queryable: boolean;
             }[];
 
-            return searchIndexes.some((index) => indexValidator(index));
+            if (!searchIndexes.some((index) => indexValidator(index))) {
+                throw new Error("Search index did not pass validation");
+            }
         },
         {
-            retries: SEARCH_INDEX_STATUS_CHECK_RETRIES,
-            retryTimeout: 100,
-            abortSignal,
-            shouldRetryOnError: (error: unknown) => {
-                return (
-                    error instanceof Error &&
-                    error.message.includes("Error connecting to Search Index Management service")
-                );
-            },
+            timeout,
+            interval,
         }
     );
 }
 
 export async function waitUntilSearchIndexIsListed(
     collection: Collection,
     searchIndex: string,
-    abortSignal?: AbortSignal
+    timeout: number = DEFAULT_WAIT_TIMEOUT,
+    interval: number = DEFAULT_RETRY_INTERVAL
 ): Promise<void> {
-    return waitUntilSearchIndexIs(collection, searchIndex, (index) => index.name === searchIndex, abortSignal);
+    return waitUntilSearchIndexIs(collection, searchIndex, (index) => index.name === searchIndex, timeout, interval);
 }
 
 export async function waitUntilSearchIndexIsQueryable(
     collection: Collection,
     searchIndex: string,
-    abortSignal?: AbortSignal
+    timeout: number = DEFAULT_WAIT_TIMEOUT,
+    interval: number = DEFAULT_RETRY_INTERVAL
 ): Promise<void> {
     return waitUntilSearchIndexIs(
         collection,
         searchIndex,
         (index) => index.name === searchIndex && index.queryable,
-        abortSignal
+        timeout,
+        interval
     );
 }

tests/integration/tools/mongodb/delete/dropIndex.test.ts

@@ -56,7 +56,7 @@ function setupForVectorSearchIndexes(integration: MongoDBIntegrationTestCase): {
 } {
     let moviesCollection: Collection;
     const indexName = "searchIdx";
-    beforeEach(async ({ signal }) => {
+    beforeEach(async () => {
         await integration.connectMcpClient();
         const mongoClient = integration.mongoClient();
         moviesCollection = mongoClient.db("mflix").collection("movies");
@@ -66,12 +66,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, signal);
+        await waitUntilSearchManagementServiceIsReady(moviesCollection, SEARCH_TIMEOUT);
         await moviesCollection.createSearchIndex({
             name: indexName,
             definition: { mappings: { dynamic: true } },
         });
-        await waitUntilSearchIndexIsListed(moviesCollection, indexName, signal);
+        await waitUntilSearchIndexIsListed(moviesCollection, indexName, SEARCH_TIMEOUT);
     });
 
     afterEach(async () => {

tests/integration/tools/mongodb/search/listSearchIndexes.test.ts

@@ -44,10 +44,10 @@ describeWithMongoDB(
     (integration) => {
         let fooCollection: Collection;
 
-        beforeEach(async ({ signal }) => {
+        beforeEach(async () => {
             await integration.connectMcpClient();
             fooCollection = integration.mongoClient().db("any").collection("foo");
-            await waitUntilSearchManagementServiceIsReady(fooCollection, signal);
+            await waitUntilSearchManagementServiceIsReady(fooCollection, SEARCH_TIMEOUT);
         });
 
         describe("when the collection does not exist", () => {
@@ -77,9 +77,9 @@ describeWithMongoDB(
         });
 
         describe("when there are indexes", () => {
-            beforeEach(async ({ signal }) => {
+            beforeEach(async () => {
                 await fooCollection.insertOne({ field1: "yay" });
-                await waitUntilSearchManagementServiceIsReady(fooCollection, signal);
+                await waitUntilSearchManagementServiceIsReady(fooCollection, SEARCH_TIMEOUT);
                 await fooCollection.createSearchIndexes([{ definition: { mappings: { dynamic: true } } }]);
             });
 
@@ -99,8 +99,8 @@ describeWithMongoDB(
             it(
                 "returns the list of existing indexes and detects if they are queryable",
                 { timeout: SEARCH_TIMEOUT },
-                async ({ signal }) => {
-                    await waitUntilSearchIndexIsQueryable(fooCollection, "default", signal);
+                async () => {
+                    await waitUntilSearchIndexIsQueryable(fooCollection, "default", SEARCH_TIMEOUT);
 
                     const response = await integration.mcpClient().callTool({
                         name: "list-search-indexes",