Merge pull request #31 from upstash/dx-914-expand-vectorjs-tests-with…

…-actual-usage-convention

DX-914: Extend Tests with Index Client

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

@@ -1,6 +1,7 @@
 import { afterAll, describe, expect, test } from "bun:test";
 import { DeleteCommand, UpsertCommand } from "@commands/index";
 import { newHttpClient, randomID, range, resetIndexes } from "@utils/test-utils";
+import { Index } from "@utils/test-utils";
 
 const client = newHttpClient();
 
@@ -47,3 +48,40 @@ describe("DELETE", () => {
     });
   });
 });
+
+describe("DELETE with Index Client", () => {
+  const index = new Index({
+    token: process.env.UPSTASH_VECTOR_REST_TOKEN!,
+    url: process.env.UPSTASH_VECTOR_REST_URL!,
+  });
+  afterAll(async () => {
+    await index.reset();
+  });
+
+  test("should delete single record succesfully", async () => {
+    const initialVector = range(0, 384);
+    const id = randomID();
+
+    index.upsert({ id, vector: initialVector });
+
+    const deletionResult = await index.delete(id);
+
+    expect(deletionResult).toEqual({
+      deleted: 1
+    });
+  });
+
+  test("should delete array of records successfully", async () => {
+    const initialVector = range(0, 384);
+    const idsToUpsert = [randomID(), randomID(), randomID()];
+
+    const upsertPromises = idsToUpsert.map((id) => index.upsert({ id, vector: initialVector }));
+
+    await Promise.all(upsertPromises);
+
+    const deletionResult = await index.delete(idsToUpsert);
+    expect(deletionResult).toEqual({
+      deleted: 3,
+    })
+  });
+});

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

@@ -55,20 +55,30 @@ describe("FETCH", () => {
 
     expect(res).toEqual([mockData]);
   });
+});
 
-  test("should fetch succesfully by index.fetch", async () => {
-    const index = new Index({
-      url: process.env.UPSTASH_VECTOR_REST_URL!,
-      token: process.env.UPSTASH_VECTOR_REST_TOKEN!,
-    });
+describe("FETCH with Index Client", () => {
+  const index = new Index({
+    token: process.env.UPSTASH_VECTOR_REST_TOKEN!,
+    url: process.env.UPSTASH_VECTOR_REST_URL!,
+  });
 
-    const randomFetch = await index.fetch([randomID()], {
-      includeMetadata: true,
-      namespace: "test",
-    });
+  test("should fetch array of records by IDs succesfully", async () => {
+    const randomizedData = new Array(20)
+      .fill("")
+      .map(() => ({ id: randomID(), vector: range(0, 384) }));
 
-    expect(randomFetch).toEqual([null]);
+    await index.upsert(randomizedData);
+
+    await awaitUntilIndexed(index);
 
+    const IDs = randomizedData.map((x) => x.id);
+    const res = await index.fetch(IDs, { includeVectors: true });
+
+    expect(res).toEqual(randomizedData);
+  });
+
+  test("should fetch single record by ID", async () => {
     const mockData = {
       id: randomID(),
       vector: range(0, 384),
@@ -87,4 +97,13 @@ describe("FETCH", () => {
 
     expect(fetchWithID).toEqual([mockData]);
   });
+
+  test("should return null when ID does not exist", async () => {
+    const randomFetch = await index.fetch([randomID()], {
+      includeMetadata: true,
+      namespace: "test",
+    });
+
+    expect(randomFetch).toEqual([null]);
+  });
 });

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

@@ -1,11 +1,25 @@
 import { afterAll, describe, expect, test } from "bun:test";
 import { QueryCommand, UpsertCommand } from "@commands/index";
-import { awaitUntilIndexed, newHttpClient, range, resetIndexes } from "@utils/test-utils";
+import { awaitUntilIndexed, newHttpClient, randomID, range } from "@utils/test-utils";
+import { Index } from "@utils/test-utils";
 
 const client = newHttpClient();
 
 describe("QUERY", () => {
-  afterAll(async () => await resetIndexes());
+  const index = new Index({
+    token: process.env.UPSTASH_VECTOR_REST_TOKEN!,
+    url: process.env.UPSTASH_VECTOR_REST_URL!,
+  });
+
+  const embeddingIndex = new Index({
+    token: process.env.EMBEDDING_UPSTASH_VECTOR_REST_TOKEN!,
+    url: process.env.EMBEDDING_UPSTASH_VECTOR_REST_URL!,
+  });
+
+  afterAll(async () => {
+    await index.reset();
+    await embeddingIndex.reset();
+  });
   test("should query records successfully", async () => {
     const initialVector = range(0, 384);
     const initialData = { id: 33, vector: initialVector };
@@ -180,3 +194,117 @@ describe("QUERY", () => {
     { timeout: 20000 }
   );
 });
+
+describe("with Index Client", () => {
+  const index = new Index({
+    token: process.env.UPSTASH_VECTOR_REST_TOKEN!,
+    url: process.env.UPSTASH_VECTOR_REST_URL!,
+  });
+  const embeddingIndex = new Index({
+    token: process.env.EMBEDDING_UPSTASH_VECTOR_REST_TOKEN!,
+    url: process.env.EMBEDDING_UPSTASH_VECTOR_REST_URL!,
+  });
+
+  afterAll(async () => {
+    await index.reset();
+    await embeddingIndex.reset();
+  });
+
+  test("should query records successfully", async () => {
+    const ID = randomID();
+    const initialVector = range(0, 384);
+    const initialData = { id: ID, vector: initialVector };
+    await index.upsert(initialData);
+
+    await awaitUntilIndexed(index);
+
+    const res = await index.query<{ hello: "World" }>({
+      includeVectors: true,
+      vector: initialVector,
+      topK: 1,
+    });
+
+    expect(res).toEqual([
+      {
+        id: ID,
+        score: 1,
+        vector: initialVector,
+      },
+    ]);
+  });
+
+  test(
+    "should query with plain text successfully",
+    async () => {
+      await embeddingIndex.upsert([
+        {
+          id: "hello-world",
+          data: "with-index-plain-text-query-test",
+          metadata: { upstash: "test" },
+        },
+      ]);
+
+      await awaitUntilIndexed(embeddingIndex);
+
+      const res = await embeddingIndex.query({
+        data: "with-index-plain-text-query-test",
+        topK: 1,
+        includeVectors: true,
+        includeMetadata: true,
+      });
+
+      expect(res[0].metadata).toEqual({ upstash: "test" });
+    },
+    { timeout: 20000 }
+  );
+
+  test("should narrow down the query results with filter", async () => {
+    const ID = randomID();
+    const initialVector = range(0, 384);
+    const initialData = [
+      {
+        id: `1-${ID}`,
+        vector: initialVector,
+        metadata: {
+          animal: "elephant",
+          tags: ["mammal"],
+          diet: "herbivore",
+        },
+      },
+      {
+        id: `2-${ID}`,
+        vector: initialVector,
+        metadata: {
+          animal: "tiger",
+          tags: ["mammal"],
+          diet: "carnivore",
+        },
+      },
+    ];
+
+    await index.upsert(initialData);
+
+    await awaitUntilIndexed(index);
+
+    const res = await index.query<{
+      animal: string;
+      tags: string[];
+      diet: string;
+    }>({
+      vector: initialVector,
+      topK: 1,
+      filter: "tags[0] = 'mammal' AND diet = 'carnivore'",
+      includeVectors: true,
+      includeMetadata: true,
+    });
+
+    expect(res).toEqual([
+      {
+        id: `2-${ID}`,
+        score: 1,
+        vector: initialVector,
+        metadata: { animal: "tiger", tags: ["mammal"], diet: "carnivore" },
+      },
+    ]);
+  });
+});

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

@@ -1,20 +1,29 @@
 import { afterAll, describe, expect, test } from "bun:test";
 import { RangeCommand, UpsertCommand } from "@commands/index";
-import { newHttpClient, randomID, range, resetIndexes } from "@utils/test-utils";
+import {
+  Index,
+  awaitUntilIndexed,
+  newHttpClient,
+  randomID,
+  range,
+  resetIndexes,
+} from "@utils/test-utils";
 
 const client = newHttpClient();
 
 describe("RANGE", () => {
   afterAll(async () => await resetIndexes());
 
-  test("should query records successfully", async () => {
+  test("should paginate records successfully", async () => {
     const randomizedData = new Array(20)
       .fill("")
       .map(() => ({ id: randomID(), vector: range(0, 384) }));
 
     const payloads = randomizedData.map((data) => new UpsertCommand(data).exec(client));
     await Promise.all(payloads);
 
+    await awaitUntilIndexed(client);
+
     const res = await new RangeCommand({
       cursor: 0,
       limit: 5,
@@ -23,3 +32,31 @@ describe("RANGE", () => {
     expect(res.nextCursor).toBe("5");
   });
 });
+
+describe("RANGE with Index Client", () => {
+  const index = new Index({
+    token: process.env.UPSTASH_VECTOR_REST_TOKEN!,
+    url: process.env.UPSTASH_VECTOR_REST_URL!,
+  });
+
+  afterAll(async () => {
+    await index.reset();
+  });
+  test("should paginate records successfully", async () => {
+    const randomizedData = new Array(20)
+      .fill("")
+      .map(() => ({ id: randomID(), vector: range(0, 384) }));
+
+    await index.upsert(randomizedData);
+
+    await awaitUntilIndexed(index);
+
+    const res = await index.range({
+      cursor: 0,
+      limit: 5,
+      includeVectors: true,
+    });
+
+    expect(res.nextCursor).toBe("5");
+  });
+});

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

@@ -1,6 +1,6 @@
 import { afterAll, describe, expect, test } from "bun:test";
 import { FetchCommand, UpdateCommand, UpsertCommand } from "@commands/index";
-import { awaitUntilIndexed, newHttpClient, range, resetIndexes } from "@utils/test-utils";
+import { Index, awaitUntilIndexed, newHttpClient, range, resetIndexes } from "@utils/test-utils";
 
 const client = newHttpClient();
 
@@ -31,3 +31,36 @@ describe("UPDATE", () => {
     expect(fetchData[0]?.metadata?.upstash).toBe("test-update");
   });
 });
+
+describe("UPDATE with Index Client", () => {
+  const index = new Index({
+    token: process.env.UPSTASH_VECTOR_REST_TOKEN!,
+    url: process.env.UPSTASH_VECTOR_REST_URL!,
+  });
+  afterAll(async () => {
+    await index.reset();
+  });
+
+  test("should update vector metadata", async () => {
+    await index.upsert({
+      id: 1,
+      vector: range(0, 384),
+      metadata: { upstash: "test-simple" },
+    });
+
+    await awaitUntilIndexed(index);
+
+    const res = await index.update({
+      id: 1,
+      metadata: { upstash: "test-update" },
+    });
+
+    expect(res).toEqual({ updated: 1 });
+
+    await awaitUntilIndexed(client, 5000);
+
+    const fetchData = await index.fetch(["1"], { includeMetadata: true });
+
+    expect(fetchData[0]?.metadata?.upstash).toBe("test-update");
+  });
+});