gadget/gadget-code/tests/vector-store.test.ts

import { describe, it, expect, vi, beforeEach } from "vitest";

const mockEmbeddings = vi.hoisted(() => vi.fn());
const mockFindById = vi.hoisted(() =>
  vi.fn().mockResolvedValue({
    _id: "test-provider-id",
    name: "Test Provider",
    apiType: "ollama",
    baseUrl: "http://test:11434",
    apiKey: "",
  }),
);

// Create mock methods shared between QdrantClient mock and test code
const mockQdrantMethods = vi.hoisted(() => ({
  getCollections: vi.fn(),
  createCollection: vi.fn(),
  getCollection: vi.fn(),
  search: vi.fn(),
  upsert: vi.fn(),
  delete: vi.fn(),
}));

vi.mock("../src/config/env.js", () => ({
  default: {
    NODE_ENV: "test",
    INSTALL_DIR: "/tmp",
    timezone: "UTC",
    pkg: { name: "test", version: "0.0.0" },
    site: {},
    ai: { ollama: { apiUrl: "http://test:11434", apiKey: "" } },
    auth: { jwtSecret: "test-secret" },
    session: {
      secret: "test-secret",
      trustProxy: false,
      cookie: { secure: false, sameSite: false },
    },
    google: { cse: { apiKey: "", engineId: "" } },
    mongodb: { host: "localhost:27017", database: "test" },
    qdrant: {
      host: "localhost",
      port: 6333,
      collection: "test-collection",
      providerId: "test-provider-id",
      embeddingModel: "test-model",
      vectorSize: 1024,
    },
    redis: {
      host: "localhost",
      port: 6379,
      password: "",
      keyPrefix: "test:",
      lazyConnect: true,
    },
    minio: {
      endpoint: "localhost",
      port: 9000,
      useSsl: false,
      accessKey: "",
      secretKey: "",
      buckets: { uploads: "", images: "", videos: "", audios: "" },
    },
    user: { passwordSalt: "test-salt" },
    https: { enabled: false, address: "localhost", port: 3443 },
    socket: { maxHttpBufferSize: 1048576 },
    frontend: { port: 5173 },
    email: { enabled: false, smtp: {}, contact: {} },
    log: {
      https: { enabled: false },
      console: { enabled: false },
      file: { enabled: false },
    },
  },
}));

vi.mock("../src/models/ai-provider.js", () => ({
  default: { findById: mockFindById },
}));

vi.mock("@gadget/ai", () => ({
  createAiApi: vi.fn().mockReturnValue({
    embeddings: mockEmbeddings,
  }),
}));

vi.mock("@langchain/textsplitters", () => ({
  RecursiveCharacterTextSplitter: vi.fn(function () {
    return {
      splitText: vi.fn().mockResolvedValue(["chunk1"]),
    };
  }),
}));

// QdrantClient mock constructor that returns the shared mock methods
vi.mock("@qdrant/js-client-rest", () => ({
  QdrantClient: vi.fn(function () {
    return mockQdrantMethods;
  }),
}));

import VectorStoreService from "../src/services/vector-store.js";

const svc = VectorStoreService as unknown as {
  _initialized: boolean;
  _dimensionMismatch: boolean;
  client: typeof mockQdrantMethods;
  aiApi: {
    embeddings: ReturnType<typeof vi.fn>;
  };
};

function setMockDefaults() {
  mockQdrantMethods.getCollections.mockResolvedValue({ collections: [] });
  mockQdrantMethods.createCollection.mockResolvedValue(undefined);
  mockQdrantMethods.getCollection.mockResolvedValue({});
  mockQdrantMethods.search.mockResolvedValue([]);
  mockQdrantMethods.upsert.mockResolvedValue(undefined);
  mockQdrantMethods.delete.mockResolvedValue(undefined);
}

function ensureInstanceSetup() {
  // For tests that don't call start(), manually assign client/aiApi
  // since start() normally does this via new QdrantClient() + createAiApi()
  if (!svc.client) {
    (svc as Record<string, unknown>).client = mockQdrantMethods;
  }
  if (!svc.aiApi) {
    (svc as Record<string, unknown>).aiApi = { embeddings: mockEmbeddings };
  }
}

describe("VectorStoreService", () => {
  beforeEach(() => {
    vi.clearAllMocks();
    setMockDefaults();

    // Reset internal state on the instance
    svc._initialized = false;
    svc._dimensionMismatch = false;
  });

  describe("search", () => {
    beforeEach(() => {
      ensureInstanceSetup();
    });

    it("throws when service is not initialized", async () => {
      svc._initialized = false;
      svc._dimensionMismatch = false;

      await expect(VectorStoreService.search("test query")).rejects.toThrow(
        "VectorStoreService is not initialized",
      );
    });

    it("throws when dimension mismatch flag is set", async () => {
      svc._initialized = true;
      svc._dimensionMismatch = true;

      await expect(VectorStoreService.search("test query")).rejects.toThrow(
        "Vector dimension mismatch: the Qdrant collection dimensions do not match the configured vectorSize (1024)",
      );
    });

    it("throws when query embedding dimensions mismatch config", async () => {
      svc._initialized = true;
      svc._dimensionMismatch = false;

      mockEmbeddings.mockResolvedValue({
        embedding: new Array(512).fill(0.1),
        model: "test-model",
      });

      await expect(VectorStoreService.search("test query")).rejects.toThrow(
        "Embedding dimension mismatch: model produced 512 dimensions, but collection expects 1024",
      );
    });

    it("passes vectorSize dimensions to the AI API when embedding", async () => {
      svc._initialized = true;
      svc._dimensionMismatch = false;

      mockQdrantMethods.search.mockResolvedValue([]);
      mockEmbeddings.mockResolvedValue({
        embedding: new Array(1024).fill(0.1),
        model: "test-model",
      });

      await VectorStoreService.search("test query");

      expect(mockEmbeddings).toHaveBeenCalledWith(
        "test-model",
        "test query",
        1024,
      );
    });

    it("returns hydrated search results", async () => {
      svc._initialized = true;
      svc._dimensionMismatch = false;

      mockEmbeddings.mockResolvedValue({
        embedding: new Array(1024).fill(0.1),
        model: "test-model",
      });

      mockQdrantMethods.search.mockResolvedValue([
        {
          id: "point-1",
          score: 0.95,
          payload: {
            content: "some content",
            userId: "user-1",
            projectId: "proj-1",
            sessionId: "session-1",
            turnId: "turn-1",
            role: "user",
            createdAt: "2026-01-01T00:00:00.000Z",
          },
        },
      ]);

      const results = await VectorStoreService.search("test query", undefined, 5);
      expect(results).toHaveLength(1);
      expect(results[0]).toMatchObject({
        id: "point-1",
        content: "some content",
        score: 0.95,
        userId: "user-1",
      });
    });
  });

  describe("start / ensureCollection", () => {
    it("creates collection when it does not exist", async () => {
      mockQdrantMethods.getCollections.mockResolvedValue({
        collections: [],
      });
      mockQdrantMethods.createCollection.mockResolvedValue(undefined);
      mockEmbeddings.mockResolvedValue({
        embedding: new Array(1024).fill(0.1),
        model: "test-model",
      });

      await VectorStoreService.start();

      expect(mockQdrantMethods.createCollection).toHaveBeenCalledWith(
        "test-collection",
        { vectors: { size: 1024, distance: "Cosine" } },
      );
      expect(svc._initialized).toBe(true);
      expect(svc._dimensionMismatch).toBe(false);
    });

    it("detects dimension mismatch when existing collection has wrong size", async () => {
      mockQdrantMethods.getCollections.mockResolvedValue({
        collections: [{ name: "test-collection" }],
      });
      mockQdrantMethods.getCollection.mockResolvedValue({
        config: {
          params: {
            vectors: { size: 768, distance: "Cosine" },
          },
        },
      });
      mockEmbeddings.mockResolvedValue({
        embedding: new Array(1024).fill(0.1),
        model: "test-model",
      });

      await VectorStoreService.start();

      expect(svc._initialized).toBe(true);
      expect(svc._dimensionMismatch).toBe(true);
    });

    it("sets dimension mismatch when embedding model produces wrong dimensions", async () => {
      mockQdrantMethods.getCollections.mockResolvedValue({
        collections: [],
      });
      mockQdrantMethods.createCollection.mockResolvedValue(undefined);
      mockEmbeddings.mockResolvedValue({
        embedding: new Array(512).fill(0.1),
        model: "test-model",
      });

      await VectorStoreService.start();

      expect(svc._initialized).toBe(true);
      expect(svc._dimensionMismatch).toBe(true);
    });

    it("starts cleanly when everything matches", async () => {
      mockQdrantMethods.getCollections.mockResolvedValue({
        collections: [],
      });
      mockQdrantMethods.createCollection.mockResolvedValue(undefined);
      mockEmbeddings.mockResolvedValue({
        embedding: new Array(1024).fill(0.1),
        model: "test-model",
      });

      await VectorStoreService.start();

      expect(svc._dimensionMismatch).toBe(false);
      expect(svc._initialized).toBe(true);
    });

    it("skips start when no providerId is configured", async () => {
      const env = await import("../src/config/env.js");
      const originalProviderId = env.default.qdrant.providerId;
      env.default.qdrant.providerId = "";

      await VectorStoreService.start();
      expect(svc._initialized).toBe(false);

      env.default.qdrant.providerId = originalProviderId;
    });
  });
});