Skip to content

Commit

Permalink
Add reset command
Browse files Browse the repository at this point in the history
  • Loading branch information
ogzhanolguncu committed Jan 11, 2024
1 parent 48f4af8 commit e20f705
Show file tree
Hide file tree
Showing 9 changed files with 70 additions and 11 deletions.
6 changes: 4 additions & 2 deletions src/commands/client/delete/index.test.ts
Original file line number Diff line number Diff line change
@@ -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()];
Expand Down
7 changes: 4 additions & 3 deletions src/commands/client/fetch/index.test.ts
Original file line number Diff line number Diff line change
@@ -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("")
Expand Down
11 changes: 8 additions & 3 deletions src/commands/client/query/index.test.ts
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
34 changes: 34 additions & 0 deletions src/commands/client/reset/index.test.ts
Original file line number Diff line number Diff line change
@@ -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));
});
});
7 changes: 7 additions & 0 deletions src/commands/client/reset/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
import { Command } from "../../index";

export class ResetCommand extends Command<string> {
constructor() {
super([], "reset");
}
}
6 changes: 4 additions & 2 deletions src/commands/client/upsert/index.test.ts
Original file line number Diff line number Diff line change
@@ -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");
Expand Down
3 changes: 2 additions & 1 deletion src/commands/index.ts
Original file line number Diff line number Diff line change
@@ -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.
Expand Down
5 changes: 5 additions & 0 deletions src/utils/test-utils.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import { ResetCommand } from "../commands/client/reset";
import { HttpClient, RetryConfig } from "../http";

export const newHttpClient = (retry?: RetryConfig | undefined) => {
Expand Down Expand Up @@ -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());
2 changes: 2 additions & 0 deletions src/vector.ts
Original file line number Diff line number Diff line change
@@ -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";

Expand Down Expand Up @@ -32,4 +33,5 @@ export class Index {
query = (args: CommandArgs<typeof QueryCommand>) => new QueryCommand(args).exec(this.client);
upsert = (args: CommandArgs<typeof UpsertCommand>) => new UpsertCommand(args).exec(this.client);
fetch = (args: CommandArgs<typeof FetchCommand>) => new FetchCommand(args).exec(this.client);
reset = () => new ResetCommand().exec(this.client);
}

0 comments on commit e20f705

Please sign in to comment.