obsidian_ollama/tests/vault-vector-store.test.ts

import { ChromaClient } from 'chromadb';
import { VaultIndexConfig } from '../src/types';

// Mock ChromaDB module
jest.mock('chromadb', () => ({
  ChromaClient: jest.fn().mockImplementation(() => {
    return {
      getOrCreateCollection: jest.fn().mockResolvedValue({
        query: jest.fn(),
        upsert: jest.fn(),
        delete: jest.fn(),
        count: jest.fn().mockResolvedValue(5),
      }),
      deleteCollection: jest.fn(),
    };
  }),
}));

import { VaultVectorStore } from '../src/vault-vector-store';

describe('VaultVectorStore', () => {
  const mockOllamaUrl = 'http://localhost:11434';
  const mockConfig: VaultIndexConfig = {
    enabled: true,
    similarityThreshold: 0.75,
    collectionName: 'test_vault_index',
    embeddingModel: 'nomic-embed-text',
    chromaURL: 'http://localhost:8000',
  };

  let store: VaultVectorStore;
  let mockChromaClient: any;
  let mockCollection: any;

  beforeEach(async () => {
    jest.clearAllMocks();
    global.fetch = jest.fn().mockResolvedValue({
      ok: true,
      json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }),
    });

    store = new VaultVectorStore(mockOllamaUrl, mockConfig);
    await store.initialize();

    mockChromaClient = (ChromaClient as jest.Mock).mock.results[0].value;
    mockCollection = await mockChromaClient.getOrCreateCollection.mock.results[0].value;
  });

  describe('constructor', () => {
    it('should initialize with provided config', () => {
      expect(store).toBeInstanceOf(VaultVectorStore);
    });
  });

  describe('initialize', () => {
    it('should initialize the collection', async () => {
      expect(mockChromaClient.getOrCreateCollection).toHaveBeenCalledWith({
        name: mockConfig.collectionName,
        metadata: { 'hnsw:space': 'cosine' },
      });
    });

    it('should not initialize when disabled', async () => {
      jest.clearAllMocks();
      const disabledConfig: VaultIndexConfig = { ...mockConfig, enabled: false };
      const disabledStore = new VaultVectorStore(mockOllamaUrl, disabledConfig);
      await disabledStore.initialize();
      expect(mockChromaClient.getOrCreateCollection).not.toHaveBeenCalled();
    });
  });

  describe('indexFile', () => {
    it('should upsert a file into the collection', async () => {
      const mockFile = {
        basename: 'test.md',
        path: 'test.md',
        extension: 'md',
      } as any;

      await store.indexFile(mockFile, '# Test\n\nThis is test content.');

      expect(mockCollection.upsert).toHaveBeenCalled();
      const upsertCall = mockCollection.upsert.mock.calls[0][0];
      expect(upsertCall.ids).toContain('test.md');
      expect(upsertCall.metadatas[0].title).toBe('test');
    });

    it('should delete file from index when content is empty', async () => {
      const mockFile = {
        basename: 'empty.md',
        path: 'empty.md',
        extension: 'md',
      } as any;

      await store.indexFile(mockFile, '   ');
      expect(mockCollection.delete).toHaveBeenCalledWith({ ids: ['empty.md'] });
    });

    it('should not index when collection is null', async () => {
      jest.clearAllMocks();
      const disabledConfig: VaultIndexConfig = { ...mockConfig, enabled: false };
      const disabledStore = new VaultVectorStore(mockOllamaUrl, disabledConfig);
      await disabledStore.initialize();

      const mockFile = { basename: 'test.md', path: 'test.md' } as any;
      await disabledStore.indexFile(mockFile, 'content');
      expect(mockCollection.upsert).not.toHaveBeenCalled();
    });
  });

  describe('deleteFile', () => {
    it('should delete a file from the collection', async () => {
      await store.deleteFile('test.md');
      expect(mockCollection.delete).toHaveBeenCalledWith({ ids: ['test.md'] });
    });
  });

  describe('search', () => {
    it('should return empty array when disabled', async () => {
      jest.clearAllMocks();
      const disabledConfig: VaultIndexConfig = { ...mockConfig, enabled: false };
      const disabledStore = new VaultVectorStore(mockOllamaUrl, disabledConfig);
      await disabledStore.initialize();

      const results = await disabledStore.search('test query', 3);
      expect(results).toEqual([]);
      expect(mockCollection.query).not.toHaveBeenCalled();
    });

    it('should return semantic search results', async () => {
      mockCollection.query.mockResolvedValue({
        ids: [['file1.md', 'file2.md']],
        documents: [['Doc 1 content', 'Doc 2 content']],
        distances: [[0.1, 0.15]],
        metadatas: [
          [
            { path: 'file1.md', title: 'File 1' },
            { path: 'file2.md', title: 'File 2' },
          ],
        ],
      });

      const results = await store.search('test query', 2);

      expect(results).toHaveLength(2);
      expect(results[0].path).toBe('file1.md');
      expect(results[0].title).toBe('File 1');
      expect(results[0].score).toBe(0.9); // 1 - 0.1
      expect(results[1].score).toBe(0.85); // 1 - 0.15
    });

    it('should filter results below similarity threshold', async () => {
      mockCollection.query.mockResolvedValue({
        ids: [['file1.md', 'file2.md']],
        documents: [['Doc 1', 'Doc 2']],
        distances: [[0.1, 0.5]], // scores: 0.9 and 0.5; threshold is 0.75
        metadatas: [
          [
            { path: 'file1.md', title: 'File 1' },
            { path: 'file2.md', title: 'File 2' },
          ],
        ],
      });

      const results = await store.search('test', 2);
      expect(results).toHaveLength(1);
      expect(results[0].path).toBe('file1.md');
    });

    it('should return empty array when no results', async () => {
      mockCollection.query.mockResolvedValue({
        ids: [[]],
        documents: [[]],
        distances: [[]],
        metadatas: [[]],
      });

      const results = await store.search('test', 3);
      expect(results).toEqual([]);
    });

    it('should return empty array on query failure', async () => {
      mockCollection.query.mockRejectedValue(new Error('Query failed'));

      const results = await store.search('test', 3);
      expect(results).toEqual([]);
    });
  });

  describe('clearIndex', () => {
    it('should delete the collection', async () => {
      await store.clearIndex();
      expect(mockChromaClient.deleteCollection).toHaveBeenCalledWith({
        name: mockConfig.collectionName,
      });
    });

    it('should not clear when disabled', async () => {
      jest.clearAllMocks();
      const disabledConfig: VaultIndexConfig = { ...mockConfig, enabled: false };
      const disabledStore = new VaultVectorStore(mockOllamaUrl, disabledConfig);
      await disabledStore.clearIndex();
      expect(mockChromaClient.deleteCollection).not.toHaveBeenCalled();
    });
  });

  describe('getIndexedCount', () => {
    it('should return the collection count', async () => {
      const count = await store.getIndexedCount();
      expect(count).toBe(5);
      expect(mockCollection.count).toHaveBeenCalled();
    });

    it('should return 0 when disabled', async () => {
      jest.clearAllMocks();
      const disabledConfig: VaultIndexConfig = { ...mockConfig, enabled: false };
      const disabledStore = new VaultVectorStore(mockOllamaUrl, disabledConfig);
      const count = await disabledStore.getIndexedCount();
      expect(count).toBe(0);
    });
  });
});