import { beforeEach, describe, expect, test } from "@jest/globals";
import { ChromaClient } from "../src/ChromaClient";
import { ChromaNotFoundError } from "../src/Errors";

describe("upsert records", () => {
  // connects to the unauthenticated chroma instance started in
  // the global jest setup file.
  const client = new ChromaClient({
    path: process.env.DEFAULT_CHROMA_INSTANCE_URL,
  });

  beforeEach(async () => {
    await client.reset();
  });

  test("it should upsert embeddings to a collection", async () => {
    const collection = await client.createCollection({ name: "test" });
    const ids = ["test1", "test2"];
    const embeddings = [
      [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
      [10, 9, 8, 7, 6, 5, 4, 3, 2, 1],
    ];
    await collection.add({ ids, embeddings });
    const count = await collection.count();
    expect(count).toBe(2);

    const ids2 = ["test2", "test3"];
    const embeddings2 = [
      [1, 2, 3, 4, 5, 6, 7, 8, 9, 15],
      [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
    ];

    await collection.upsert({
      ids: ids2,
      embeddings: embeddings2,
    });

    const count2 = await collection.count();
    expect(count2).toBe(3);
  });

  test("should error on non existing collection", async () => {
    const collection = await client.createCollection({ name: "test" });
    await client.deleteCollection({ name: "test" });
    await expect(async () => {
      await collection.upsert({
        ids: ["test1"],
        embeddings: [[1, 2, 3, 4, 5, 6, 7, 8, 9, 11]],
        metadatas: [{ test: "meta1" }],
        documents: ["doc1"],
      });
    }).rejects.toThrow(ChromaNotFoundError);
  });
});
