obsidian_ollama/tests/semantic-cache.test.ts

// tests/semantic-cache.test.ts

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

// Mock ChromaDB module
jest.mock('chromadb', () => ({
  ChromaClient: jest.fn().mockImplementation(() => {
    return {
      getOrCreateCollection: jest.fn().mockResolvedValue({
        query: jest.fn(),
        add: jest.fn(),
        upsert: jest.fn(),
        reset: jest.fn(),
      }),
      deleteCollection: jest.fn(),
      reset: jest.fn(),
    };
  }),
  IncludeEnum: {
    Documents: 'documents',
    Embeddings: 'embeddings',
    Metadatas: 'metadatas',
    Distances: 'distances',
  },
}));

import { SemanticCacheService } from '../src/semantic-cache';

describe('SemanticCacheService', () => {
  const mockOllamaUrl = 'http://localhost:11434';
  const mockCacheConfig: CacheConfig = {
    enabled: true,
    similarityThreshold: 0.85,
    collectionName: 'test_cache',
    embeddingModel: 'nomic-embed-text',
    chromaURL: 'http://localhost:8000',
  };

  let cacheService: SemanticCacheService;
  let mockChromaClient: any;
  let mockCollection: any;

  beforeEach(async () => {
    // Reset all mocks
    jest.clearAllMocks();

    // Create a fresh instance for each test
    cacheService = new SemanticCacheService(mockOllamaUrl, mockCacheConfig);
    await cacheService.initialize();
    global.fetch = jest.fn().mockResolvedValue({
      ok: true,
      json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }),
    });

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

  describe('constructor', () => {
    it('should initialize with correct configuration', () => {
      expect(mockChromaClient).toBeDefined();
      expect(mockChromaClient.getOrCreateCollection).toHaveBeenCalledWith({
        name: mockCacheConfig.collectionName,
        metadata: { 'hnsw:space': 'cosine' },
      });
    });
  });

  describe('initialize', () => {
    it('should initialize the cache collection', async () => {
      await cacheService.initialize();

      expect(mockChromaClient.getOrCreateCollection).toHaveBeenCalledWith({
        name: mockCacheConfig.collectionName,
        metadata: { 'hnsw:space': 'cosine' },
      });
    });

    it('should not initialize when cache is disabled', async () => {
      jest.clearAllMocks();
      const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false };
      const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig);

      await disabledCacheService.initialize();

      expect(mockChromaClient.getOrCreateCollection).not.toHaveBeenCalled();
    });
  });

  describe('getCache', () => {
    it('should return null when cache is disabled', async () => {
      jest.clearAllMocks();
      const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false };
      const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig);

      const result = await disabledCacheService.getCache('test query');

      expect(result).toBeNull();
      expect(mockCollection.query).not.toHaveBeenCalled();
    });

    it('should return null when no cache hit', async () => {
      mockCollection.query.mockResolvedValue({
        ids: [[]],
        documents: [[]],
        distances: [[]],
      });

      const result = await cacheService.getCache('test query');

      expect(result).toBeNull();
      expect(mockCollection.query).toHaveBeenCalled();
    });

    it('should return cached content when hit', async () => {
      const cachedContent = 'cached response';
      mockCollection.query.mockResolvedValue({
        ids: [['test-id']],
        documents: [[cachedContent]],
        distances: [[0.9]], // Above threshold
      });

      const result = await cacheService.getCache('test query');

      expect(result).toBe(cachedContent);
      expect(mockCollection.query).toHaveBeenCalled();
    });
  });

  describe('setCache', () => {
    it('should not set cache when disabled', async () => {
      jest.clearAllMocks();
      const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false };
      const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig);

      await disabledCacheService.setCache('test query', 'test response');

      expect(mockCollection.upsert).not.toHaveBeenCalled();
    });

    it('should add content to cache', async () => {
      await cacheService.setCache('test query', 'test response');

      expect(mockCollection.upsert).toHaveBeenCalled();
    });
  });

  describe('clearCache', () => {
    it('should not clear when cache is disabled', async () => {
      jest.clearAllMocks();
      const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false };
      const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig);

      await disabledCacheService.clearCache();

      expect(mockChromaClient.reset).not.toHaveBeenCalled();
    });

    it('should clear the cache via the ChromaClient', async () => {
      await cacheService.clearCache();

      expect(mockChromaClient.reset).toHaveBeenCalled();
    });
  });
});