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); }); }); });