Merge pull request #21 from upstash/DX-829

test: simplify tests by removing second vector for testing embeddings
upstash · Apr 15, 2024 · 801870c · 801870c
2 parents 184aff1 + f483476
commit 801870c
Show file tree

Hide file tree

Showing 9 changed files with 41 additions and 46 deletions.
diff --git a/package.json b/package.json
@@ -33,7 +33,7 @@
   ],
   "license": "MIT",
   "scripts": {
-    "test": "bun test src --coverage --bail --coverageSkipTestFiles=[test-utils.ts] && vitest run --typecheck",
+    "test": "bun test src --coverage --bail --coverageSkipTestFiles=[test-utils.ts] --timeout 20000 && vitest run --typecheck",
     "fmt": "bunx biome check --apply ./src",
     "build": "tsup",
     "prepare": "husky install"

diff --git a/src/commands/client/delete/index.test.ts b/src/commands/client/delete/index.test.ts
@@ -1,14 +1,14 @@
 import { afterAll, describe, expect, test } from "bun:test";
 import { DeleteCommand, UpsertCommand } from "@commands/index";
-import { newHttpClient, randomID, resetIndexes } from "@utils/test-utils";
+import { newHttpClient, randomID, range, resetIndexes } from "@utils/test-utils";
 
 const client = newHttpClient();
 
 describe("DELETE", () => {
   afterAll(async () => await resetIndexes());
 
   test("should delete record(s) successfully", async () => {
-    const initialVector = [6.6, 7.7];
+    const initialVector = range(0, 384);
     const idsToUpsert = [randomID(), randomID(), randomID()];
 
     const upsertPromises = idsToUpsert.map((id) =>
@@ -21,7 +21,7 @@ describe("DELETE", () => {
   });
 
   test("deleting the same ids should throw", async () => {
-    const initialVector = [6.6, 7.7];
+    const initialVector = range(0, 384);
     const idsToUpsert = [randomID(), randomID(), randomID()];
 
     const upsertPromises = idsToUpsert.map((id) =>
@@ -37,7 +37,7 @@ describe("DELETE", () => {
   });
 
   test("should delete single item", async () => {
-    const initialVector = [6.6, 7.7];
+    const initialVector = range(0, 384);
     const id = randomID();
     await new UpsertCommand({ id, vector: initialVector }).exec(client);
 

diff --git a/src/commands/client/fetch/index.test.ts b/src/commands/client/fetch/index.test.ts
@@ -1,6 +1,6 @@
 import { afterAll, describe, expect, test } from "bun:test";
 import { FetchCommand, UpsertCommand } from "@commands/index";
-import { newHttpClient, randomFloat, randomID, resetIndexes } from "@utils/test-utils";
+import { newHttpClient, randomID, range, resetIndexes } from "@utils/test-utils";
 
 const client = newHttpClient();
 
@@ -10,7 +10,7 @@ describe("FETCH", () => {
   test("should fetch records successfully", async () => {
     const randomizedData = new Array(20)
       .fill("")
-      .map(() => ({ id: randomID(), vector: [randomFloat(), randomFloat()] }));
+      .map(() => ({ id: randomID(), vector: range(0, 384) }));
 
     const payloads = randomizedData.map((data) => new UpsertCommand(data).exec(client));
     await Promise.all(payloads);
@@ -39,7 +39,7 @@ describe("FETCH", () => {
   test("should return with metadata", async () => {
     const mockData = {
       id: randomID(),
-      vector: [randomFloat(), randomFloat()],
+      vector: range(0, 384),
       metadata: { hello: "world" },
     };
     await new UpsertCommand(mockData).exec(client);

diff --git a/src/commands/client/info/index.test.ts b/src/commands/client/info/index.test.ts
@@ -1,6 +1,6 @@
 import { afterAll, describe, expect, test } from "bun:test";
 import { UpsertCommand } from "@commands/index";
-import { newHttpClient, randomFloat, randomID, resetIndexes } from "@utils/test-utils";
+import { newHttpClient, randomID, range, resetIndexes } from "@utils/test-utils";
 import { sleep } from "bun";
 import { InfoCommand } from ".";
 
@@ -13,7 +13,7 @@ describe("INFO", () => {
     const vectorCount = 20;
     const randomizedData = new Array(vectorCount).fill("").map(() => ({
       id: randomID(),
-      vector: [randomFloat(), randomFloat()],
+      vector: range(0, 384),
     }));
 
     const payloads = randomizedData.map((data) => new UpsertCommand(data).exec(client));

diff --git a/src/commands/client/query/index.test.ts b/src/commands/client/query/index.test.ts
@@ -1,14 +1,14 @@
 import { afterAll, describe, expect, test } from "bun:test";
 import { QueryCommand, UpsertCommand } from "@commands/index";
-import { newHttpClient, resetIndexes } from "@utils/test-utils";
+import { newHttpClient, range, resetIndexes } from "@utils/test-utils";
 import { sleep } from "bun";
 
 const client = newHttpClient();
 
 describe("QUERY", () => {
   afterAll(async () => await resetIndexes());
   test("should query records successfully", async () => {
-    const initialVector = [6.6, 7.7];
+    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.
@@ -22,12 +22,12 @@ describe("QUERY", () => {
       {
         id: "33",
         score: 1,
-        vector: [6.6, 7.7],
+        vector: initialVector,
       },
     ]);
   });
   test("should query records filtered with metadata filter", async () => {
-    const initialVector = [6.6, 7.7];
+    const initialVector = range(0, 384);
     const initialData = {
       id: 34,
       vector: initialVector,
@@ -57,7 +57,7 @@ describe("QUERY", () => {
       {
         id: "34",
         score: 1,
-        vector: [6.6, 7.7],
+        vector: initialVector,
         metadata: {
           city: "Istanbul",
           population: 1546000,
@@ -67,11 +67,11 @@ describe("QUERY", () => {
     ]);
   });
   test("should narrow down the query results with filter", async () => {
-    const exampleVector = [6.6, 7.7];
+    const initialVector = range(0, 384);
     const initialData = [
       {
         id: 1,
-        vector: exampleVector,
+        vector: initialVector,
         metadata: {
           animal: "elephant",
           tags: ["mammal"],
@@ -80,7 +80,7 @@ describe("QUERY", () => {
       },
       {
         id: 2,
-        vector: exampleVector,
+        vector: initialVector,
         metadata: {
           animal: "tiger",
           tags: ["mammal"],
@@ -96,7 +96,7 @@ describe("QUERY", () => {
       tags: string[];
       diet: string;
     }>({
-      vector: exampleVector,
+      vector: initialVector,
       topK: 1,
       filter: "tags[0] = 'mammal' AND diet = 'carnivore'",
       includeVectors: true,
@@ -106,7 +106,7 @@ describe("QUERY", () => {
       {
         id: "2",
         score: 1,
-        vector: [6.6, 7.7],
+        vector: initialVector,
         metadata: { animal: "tiger", tags: ["mammal"], diet: "carnivore" },
       },
     ]);
@@ -115,25 +115,21 @@ describe("QUERY", () => {
   test(
     "should query with plain text successfully",
     async () => {
-      const embeddingClient = newHttpClient(undefined, {
-        token: process.env.EMBEDDING_UPSTASH_VECTOR_REST_TOKEN!,
-        url: process.env.EMBEDDING_UPSTASH_VECTOR_REST_URL!,
-      });
       await new UpsertCommand([
         {
           id: "hello-world",
           data: "Test1-2-3-4-5",
           metadata: { upstash: "test" },
         },
-      ]).exec(embeddingClient);
+      ]).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(5000);
       const res = await new QueryCommand({
         data: "Test1-2-3-4-5",
         topK: 1,
         includeVectors: true,
         includeMetadata: true,
-      }).exec(embeddingClient);
+      }).exec(client);
 
       expect(res[0].metadata).toEqual({ upstash: "test" });
     },
@@ -143,10 +139,6 @@ describe("QUERY", () => {
   test(
     "should query with plain text successfully",
     async () => {
-      const embeddingClient = newHttpClient(undefined, {
-        token: process.env.EMBEDDING_UPSTASH_VECTOR_REST_TOKEN!,
-        url: process.env.EMBEDDING_UPSTASH_VECTOR_REST_URL!,
-      });
       await new UpsertCommand([
         {
           id: "hello-world",
@@ -158,20 +150,15 @@ describe("QUERY", () => {
           data: "Test1-2-3-4-5-6",
           metadata: { upstash: "Monster" },
         },
-        {
-          id: "hello-world2",
-          data: "Test1-2-3-4-5",
-          metadata: { upstash: "Jar" },
-        },
-      ]).exec(embeddingClient);
+      ]).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(5000);
       const res = await new QueryCommand({
         data: "Test1-2-3-4-5",
         topK: 1,
         includeVectors: true,
         includeMetadata: true,
-      }).exec(embeddingClient);
+      }).exec(client);
 
       expect(res[0].metadata).toEqual({ upstash: "Cookie" });
     },

diff --git a/src/commands/client/range/index.test.ts b/src/commands/client/range/index.test.ts
@@ -1,6 +1,6 @@
 import { afterAll, describe, expect, test } from "bun:test";
 import { RangeCommand, UpsertCommand } from "@commands/index";
-import { newHttpClient, randomFloat, randomID, resetIndexes } from "@utils/test-utils";
+import { newHttpClient, randomID, range, resetIndexes } from "@utils/test-utils";
 
 const client = newHttpClient();
 
@@ -10,7 +10,7 @@ describe("RANGE", () => {
   test("should query records successfully", async () => {
     const randomizedData = new Array(20)
       .fill("")
-      .map(() => ({ id: randomID(), vector: [randomFloat(), randomFloat()] }));
+      .map(() => ({ id: randomID(), vector: range(0, 384) }));
 
     const payloads = randomizedData.map((data) => new UpsertCommand(data).exec(client));
     await Promise.all(payloads);

diff --git a/src/commands/client/reset/index.test.ts b/src/commands/client/reset/index.test.ts
@@ -1,14 +1,14 @@
 import { describe, expect, test } from "bun:test";
 import { FetchCommand, ResetCommand, UpsertCommand } from "@commands/index";
-import { newHttpClient, randomFloat, randomID } from "@utils/test-utils";
+import { newHttpClient, randomID, range } from "@utils/test-utils";
 
 const client = newHttpClient();
 
 describe("RESET", () => {
   test("should flush indexes successfully", async () => {
     const randomizedData = new Array(20)
       .fill("")
-      .map(() => ({ id: randomID(), vector: [randomFloat(), randomFloat()] }));
+      .map(() => ({ id: randomID(), vector: range(0, 384) }));
 
     const payloads = randomizedData.map((data) => new UpsertCommand(data).exec(client));
 

diff --git a/src/commands/client/upsert/index.test.ts b/src/commands/client/upsert/index.test.ts
@@ -1,14 +1,14 @@
 import { afterAll, describe, expect, test } from "bun:test";
 import { FetchCommand, UpsertCommand } from "@commands/index";
-import { newHttpClient, resetIndexes } from "@utils/test-utils";
+import { newHttpClient, range, resetIndexes } from "@utils/test-utils";
 
 const client = newHttpClient();
 
 describe("UPSERT", () => {
   afterAll(async () => await resetIndexes());
 
   test("should add record successfully", async () => {
-    const res = await new UpsertCommand({ id: 1, vector: [0.1, 0.2] }).exec(client);
+    const res = await new UpsertCommand({ id: 1, vector: range(0, 384) }).exec(client);
     expect(res).toEqual("Success");
   });
 
@@ -25,7 +25,7 @@ describe("UPSERT", () => {
     //@ts-ignore
     const res = await new UpsertCommand({
       id: 1,
-      vector: [0.1, 0.2],
+      vector: range(0, 384),
       metadata: { upstash: "test" },
     }).exec(client);
     expect(res).toEqual("Success");
@@ -35,12 +35,12 @@ describe("UPSERT", () => {
     const res = await new UpsertCommand([
       {
         id: "hello-world",
-        vector: [0.1, 0.2],
+        vector: range(0, 384),
         metadata: { upstash: "test" },
       },
       {
         id: "hello-world-4",
-        vector: [3, 4],
+        vector: range(0, 384),
         metadata: { upstash: "test" },
       },
     ]).exec(client);

diff --git a/src/utils/test-utils.ts b/src/utils/test-utils.ts
@@ -50,3 +50,11 @@ export function randomID(): string {
 export const randomFloat = () => parseFloat((Math.random() * 10).toFixed(1));
 
 export const resetIndexes = async () => await new ResetCommand().exec(newHttpClient());
+
+export const range = (start: number, end: number, step = 1) => {
+  const result = [];
+  for (let i = start; i < end; i += step) {
+    result.push(i);
+  }
+  return result;
+};