From e20f7058ae7b482e0425c07eeac355fb8adcfe73 Mon Sep 17 00:00:00 2001 From: ogzhanolguncu Date: Thu, 11 Jan 2024 12:23:10 +0300 Subject: [PATCH] Add reset command --- src/commands/client/delete/index.test.ts | 6 +++-- src/commands/client/fetch/index.test.ts | 7 ++--- src/commands/client/query/index.test.ts | 11 +++++--- src/commands/client/reset/index.test.ts | 34 ++++++++++++++++++++++++ src/commands/client/reset/index.ts | 7 +++++ src/commands/client/upsert/index.test.ts | 6 +++-- src/commands/index.ts | 3 ++- src/utils/test-utils.ts | 5 ++++ src/vector.ts | 2 ++ 9 files changed, 70 insertions(+), 11 deletions(-) create mode 100644 src/commands/client/reset/index.test.ts create mode 100644 src/commands/client/reset/index.ts diff --git a/src/commands/client/delete/index.test.ts b/src/commands/client/delete/index.test.ts index 805108d..ec30e29 100644 --- a/src/commands/client/delete/index.test.ts +++ b/src/commands/client/delete/index.test.ts @@ -1,11 +1,13 @@ -import { describe, expect, test } from "bun:test"; -import { newHttpClient, randomID } from "../../../utils/test-utils"; +import { afterAll, describe, expect, test } from "bun:test"; +import { newHttpClient, randomID, resetIndexes } from "../../../utils/test-utils"; import { UpsertCommand } from "../upsert"; import { DeleteCommand } from "./"; const client = newHttpClient(); describe("DELETE", () => { + afterAll(async () => await resetIndexes()); + test("should delete record(s) successfully", async () => { const initialVector = [6.6, 7.7]; const idsToUpsert = [randomID(), randomID(), randomID()]; diff --git a/src/commands/client/fetch/index.test.ts b/src/commands/client/fetch/index.test.ts index 7cd9b3e..ffa778c 100644 --- a/src/commands/client/fetch/index.test.ts +++ b/src/commands/client/fetch/index.test.ts @@ -1,12 +1,13 @@ -import { describe, expect, test } from "bun:test"; +import { afterAll, describe, expect, test } from "bun:test"; import { FetchCommand } from "."; -import { newHttpClient, randomID } from "../../../utils/test-utils"; +import { newHttpClient, randomFloat, randomID, resetIndexes } from "../../../utils/test-utils"; import { UpsertCommand } from "../upsert"; const client = newHttpClient(); -const randomFloat = () => parseFloat((Math.random() * 10).toFixed(1)); describe("FETCH", () => { + afterAll(async () => await resetIndexes()); + test("should fetch records successfully", async () => { const randomizedData = new Array(20) .fill("") diff --git a/src/commands/client/query/index.test.ts b/src/commands/client/query/index.test.ts index 270031d..63ce093 100644 --- a/src/commands/client/query/index.test.ts +++ b/src/commands/client/query/index.test.ts @@ -1,16 +1,21 @@ -import { describe, expect, test } from "bun:test"; -import { newHttpClient } from "../../../utils/test-utils"; +import { afterAll, describe, expect, test } from "bun:test"; +import { sleep } from "bun"; +import { newHttpClient, resetIndexes } from "../../../utils/test-utils"; import { UpsertCommand } from "../upsert"; import { QueryCommand } from "./index"; const client = newHttpClient(); describe("QUERY", () => { + afterAll(async () => await resetIndexes()); + test("should query records successfully", async () => { const initialVector = [6.6, 7.7]; const initialData = { id: 33, vector: initialVector }; - await new UpsertCommand(initialData).exec(client); + 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); const res = await new QueryCommand({ includeVectors: true, vector: initialVector, diff --git a/src/commands/client/reset/index.test.ts b/src/commands/client/reset/index.test.ts new file mode 100644 index 0000000..50533ff --- /dev/null +++ b/src/commands/client/reset/index.test.ts @@ -0,0 +1,34 @@ +import { describe, expect, test } from "bun:test"; +import { ResetCommand } from "."; +import { newHttpClient, randomFloat, randomID } from "../../../utils/test-utils"; +import { FetchCommand } from "../fetch"; +import { UpsertCommand } from "../upsert"; + +const client = newHttpClient(); + +describe("RESET", () => { + test("should flush indexes successfully", async () => { + const randomizedData = new Array(20) + .fill("") + .map(() => ({ id: randomID(), vector: [randomFloat(), randomFloat()] })); + + const payloads = randomizedData.map((data) => new UpsertCommand(data).exec(client)); + + await Promise.all(payloads); + + const res = await new FetchCommand({ + ids: randomizedData.map((x) => x.id), + includeVectors: true, + }).exec(client); + + expect(res).toEqual(randomizedData); + + await new ResetCommand().exec(client); + const resAfterReset = await new FetchCommand({ + ids: randomizedData.map((x) => x.id), + includeVectors: true, + }).exec(client); + + expect(resAfterReset).toEqual(new Array(20).fill(null)); + }); +}); diff --git a/src/commands/client/reset/index.ts b/src/commands/client/reset/index.ts new file mode 100644 index 0000000..db4170e --- /dev/null +++ b/src/commands/client/reset/index.ts @@ -0,0 +1,7 @@ +import { Command } from "../../index"; + +export class ResetCommand extends Command { + constructor() { + super([], "reset"); + } +} diff --git a/src/commands/client/upsert/index.test.ts b/src/commands/client/upsert/index.test.ts index bd4edac..e42b3be 100644 --- a/src/commands/client/upsert/index.test.ts +++ b/src/commands/client/upsert/index.test.ts @@ -1,10 +1,12 @@ -import { describe, expect, test } from "bun:test"; +import { afterAll, describe, expect, test } from "bun:test"; import { UpsertCommand } from "."; -import { newHttpClient } from "../../../utils/test-utils"; +import { newHttpClient, 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); expect(res).toEqual("Success"); diff --git a/src/commands/index.ts b/src/commands/index.ts index 0901a49..7547439 100644 --- a/src/commands/index.ts +++ b/src/commands/index.ts @@ -1,7 +1,8 @@ import { UpstashError } from "../error"; import { Requester } from "../http"; -const ENDPOINTS = ["upsert", "query", "delete", "fetch"] as const; +const ENDPOINTS = ["upsert", "query", "delete", "fetch", "reset"] as const; + export type EndpointVariants = (typeof ENDPOINTS)[number]; /** * TResult is the raw data returned from upstash, which may need to be transformed or parsed. diff --git a/src/utils/test-utils.ts b/src/utils/test-utils.ts index 3488eeb..f935c91 100644 --- a/src/utils/test-utils.ts +++ b/src/utils/test-utils.ts @@ -1,3 +1,4 @@ +import { ResetCommand } from "../commands/client/reset"; import { HttpClient, RetryConfig } from "../http"; export const newHttpClient = (retry?: RetryConfig | undefined) => { @@ -40,3 +41,7 @@ export function randomID(): string { } return btoa(s.join("")); } + +export const randomFloat = () => parseFloat((Math.random() * 10).toFixed(1)); + +export const resetIndexes = async () => await new ResetCommand().exec(newHttpClient()); diff --git a/src/vector.ts b/src/vector.ts index 5c046f2..37b0aed 100644 --- a/src/vector.ts +++ b/src/vector.ts @@ -1,6 +1,7 @@ import { FetchCommand } from "./commands/client"; import { DeleteCommand } from "./commands/client/delete"; import { QueryCommand } from "./commands/client/query"; +import { ResetCommand } from "./commands/client/reset"; import { UpsertCommand } from "./commands/client/upsert"; import { Requester } from "./http"; @@ -32,4 +33,5 @@ export class Index { query = (args: CommandArgs) => new QueryCommand(args).exec(this.client); upsert = (args: CommandArgs) => new UpsertCommand(args).exec(this.client); fetch = (args: CommandArgs) => new FetchCommand(args).exec(this.client); + reset = () => new ResetCommand().exec(this.client); }