Improve tests

src/commands/client/fetch/index.test.ts

@@ -1,7 +1,6 @@
 import { afterAll, describe, expect, test } from "bun:test";
 import { FetchCommand, UpsertCommand } from "@commands/index";
-import { newHttpClient, randomID, range, resetIndexes } from "@utils/test-utils";
-import { sleep } from "bun";
+import { awaitUntilIndexed, newHttpClient, randomID, range, resetIndexes } from "@utils/test-utils";
 import { Index } from "../../../../index";
 
 const client = newHttpClient();
@@ -78,7 +77,7 @@ describe("FETCH", () => {
 
     await index.upsert(mockData, { namespace: "test" });
 
-    sleep(4000);
+    await awaitUntilIndexed(index);
 
     const fetchWithID = await index.fetch([mockData.id], {
       includeMetadata: true,

src/commands/client/info/index.test.ts

@@ -1,7 +1,6 @@
 import { afterAll, describe, expect, test } from "bun:test";
 import { UpsertCommand } from "@commands/index";
-import { newHttpClient, randomID, range, resetIndexes } from "@utils/test-utils";
-import { sleep } from "bun";
+import { awaitUntilIndexed, newHttpClient, randomID, range, resetIndexes } from "@utils/test-utils";
 import { InfoCommand } from ".";
 
 const client = newHttpClient();
@@ -18,7 +17,9 @@ describe("INFO", () => {
 
     const payloads = randomizedData.map((data) => new UpsertCommand(data).exec(client));
     await Promise.all(payloads);
-    await sleep(2000);
+
+    await awaitUntilIndexed(client);
+
     const res = await new InfoCommand().exec(client);
     expect(res.vectorCount).toBeGreaterThanOrEqual(vectorCount);
   });

src/commands/client/namespace/index.test.ts

@@ -1,7 +1,6 @@
 import { afterAll, describe, expect, test } from "bun:test";
-import { range, resetIndexes } from "@utils/test-utils";
+import { awaitUntilIndexed, range, resetIndexes } from "@utils/test-utils";
 
-import { sleep } from "bun";
 import { Index } from "../../../../index";
 
 describe("NAMESPACE", () => {
@@ -27,7 +26,7 @@ describe("NAMESPACE", () => {
       metadata: { namespace: "namespace2" },
     });
 
-    sleep(10000);
+    await awaitUntilIndexed(index);
 
     const query1 = await namespace1.query({
       vector: range(0, 384),
@@ -55,7 +54,7 @@ describe("NAMESPACE", () => {
       metadata: { namespace: "test-namespace-reset" },
     });
 
-    sleep(5000);
+    await awaitUntilIndexed(index);
 
     const res = await namespace.query({
       vector: range(0, 384),
@@ -66,8 +65,6 @@ describe("NAMESPACE", () => {
 
     await namespace.reset();
 
-    sleep(5000);
-
     const res2 = await namespace.query({
       vector: range(0, 384),
       topK: 3,

src/commands/client/query/index.test.ts

@@ -1,7 +1,6 @@
 import { afterAll, describe, expect, test } from "bun:test";
 import { QueryCommand, UpsertCommand } from "@commands/index";
-import { newHttpClient, range, resetIndexes } from "@utils/test-utils";
-import { sleep } from "bun";
+import { awaitUntilIndexed, newHttpClient, range, resetIndexes } from "@utils/test-utils";
 
 const client = newHttpClient();
 
@@ -11,8 +10,8 @@ describe("QUERY", () => {
     const initialVector = range(0, 384);
     const initialData = { id: 33, vector: initialVector };
     await new UpsertCommand(initialData).exec(client);
-    //This is needed for vector index insertion to happen. When run with other tests in parallel this tends to fail without sleep. But, standalone it should work without an issue.
-    await sleep(2000);
+
+    await awaitUntilIndexed(client);
 
     const res = await new QueryCommand<{ hello: "World" }>({
       includeVectors: true,
@@ -42,8 +41,9 @@ describe("QUERY", () => {
       },
     };
     await new UpsertCommand(initialData).exec(client);
-    //This is needed for vector index insertion to happen. When run with other tests in parallel this tends to fail without sleep. But, standalone it should work without an issue.
-    await sleep(2000);
+
+    await awaitUntilIndexed(client);
+
     const res = await new QueryCommand<{
       city: string;
       population: number;
@@ -91,8 +91,9 @@ describe("QUERY", () => {
       },
     ];
     await new UpsertCommand(initialData).exec(client);
-    //This is needed for vector index insertion to happen. When run with other tests in parallel this tends to fail without sleep. But, standalone it should work without an issue.
-    await sleep(2000);
+
+    await awaitUntilIndexed(client);
+
     const res = await new QueryCommand<{
       animal: string;
       tags: string[];
@@ -129,8 +130,9 @@ describe("QUERY", () => {
           metadata: { upstash: "test" },
         },
       ]).exec(embeddingClient);
-      //   This is needed for vector index insertion to happen. When run with other tests in parallel this tends to fail without sleep. But, standalone it should work without an issue.
-      await sleep(5000);
+
+      await awaitUntilIndexed(embeddingClient);
+
       const res = await new QueryCommand({
         data: "Test1-2-3-4-5",
         topK: 1,
@@ -163,8 +165,9 @@ describe("QUERY", () => {
           metadata: { upstash: "Monster" },
         },
       ]).exec(embeddingClient);
-      //   This is needed for vector index insertion to happen. When run with other tests in parallel this tends to fail without sleep. But, standalone it should work without an issue.
-      await sleep(5000);
+
+      await awaitUntilIndexed(embeddingClient);
+
       const res = await new QueryCommand({
         data: "Test1-2-3-4-5",
         topK: 1,

src/commands/management/namespaces/delete/index.test.ts

@@ -1,7 +1,6 @@
 import { describe, expect, test } from "bun:test";
 import { DeleteNamespaceCommand, ListNamespacesCommand, UpsertCommand } from "@commands/index";
-import { newHttpClient, randomID, range } from "@utils/test-utils";
-import { sleep } from "bun";
+import { awaitUntilIndexed, newHttpClient, randomID, range } from "@utils/test-utils";
 
 const client = newHttpClient();
 
@@ -12,16 +11,14 @@ describe("NAMESPACES->DELETE", () => {
       { namespace: "test-namespace-delete" }
     ).exec(client);
 
-    sleep(2000);
+    await awaitUntilIndexed(client);
 
     const namespaces = await new ListNamespacesCommand().exec(client);
 
     expect(namespaces).toContain("test-namespace-delete");
 
     await new DeleteNamespaceCommand("test-namespace-delete").exec(client);
 
-    sleep(2000);
-
     const namespacesAfterDelete = await new ListNamespacesCommand().exec(client);
 
     expect(namespacesAfterDelete).not.toContain("test-namespace-delete");

src/commands/management/namespaces/list/index.test.ts

@@ -1,7 +1,6 @@
 import { describe, expect, test } from "bun:test";
 import { ListNamespacesCommand, UpsertCommand } from "@commands/index";
-import { newHttpClient, randomID, range } from "@utils/test-utils";
-import { sleep } from "bun";
+import { awaitUntilIndexed, newHttpClient, randomID, range } from "@utils/test-utils";
 
 const client = newHttpClient();
 
@@ -12,7 +11,7 @@ describe("NAMESPACES->LIST", () => {
       { namespace: "test-namespace-list" }
     ).exec(client);
 
-    sleep(2000);
+    await awaitUntilIndexed(client);
 
     const namespaces = await new ListNamespacesCommand().exec(client);
 

src/utils/test-utils.ts

@@ -1,5 +1,8 @@
+import { sleep } from "bun";
+import { InfoCommand } from "../commands/client/info";
 import { ResetCommand } from "../commands/client/reset";
 import { HttpClient, RetryConfig } from "../http";
+import { Index } from "../vector";
 
 export type NonArrayType<T> = T extends Array<infer U> ? U : T;
 
@@ -58,3 +61,29 @@ export const range = (start: number, end: number, step = 1) => {
   }
   return result;
 };
+
+export const awaitUntilIndexed = async (client: HttpClient | Index, timeoutMillis = 10_000) => {
+  const start = performance.now();
+
+  const getInfo = async () => {
+    if (client instanceof HttpClient) {
+      const cmd = new InfoCommand();
+      return await cmd.exec(client);
+    }
+
+    return await client.info();
+  };
+
+  do {
+    const info = await getInfo();
+    if (info.pendingVectorCount === 0) {
+      // OK, nothing more to index.
+      return;
+    }
+
+    // Not indexed yet, sleep a bit and check again if the timeout is not passed.
+    await sleep(1000);
+  } while (performance.now() < start + timeoutMillis);
+
+  throw new Error(`Indexing is not completed in ${timeoutMillis} ms.`);
+};