obsidian_ollama/tests/vectorization.test.ts

import { ContentVectorizer } from '../src/indexing-pipeline/vectorization';
import { ContentChunk } from '../src/indexing-pipeline/normalization';

// Mock fetch globally for all tests
global.fetch = jest.fn();

describe('ContentVectorizer', () => {
  let vectorizer: ContentVectorizer;
  let mockFetch: jest.Mock;

  beforeEach(() => {
    mockFetch = fetch as jest.Mock;
    mockFetch.mockClear();

    vectorizer = new ContentVectorizer(
      {
        model: 'test-model',
        ollamaUrl: 'http://localhost:11434',
      },
      mockFetch
    );
  });

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

    it('should use provided fetch function', async () => {
      mockFetch.mockResolvedValueOnce({
        ok: true,
        json: async () => ({ embedding: [1, 2, 3] }),
      });

      await vectorizer.vectorize({
        id: 'test-1',
        path: 'test.md',
        title: 'Test',
        content: 'Test content',
        tokens: ['test', 'content'],
        firstParagraph: 'First paragraph',
        headings: ['Heading 1'],
        frontmatter: {},
        wordCount: 2,
        chunkIndex: 0,
        chunkSize: 1,
      });

      expect(mockFetch).toHaveBeenCalled();
    });
  });

  describe('vectorize', () => {
    it('should generate embeddings for valid content', async () => {
      const mockEmbedding = [1, 2, 3, 4, 5];
      mockFetch.mockResolvedValueOnce({
        ok: true,
        json: async () => ({ embedding: mockEmbedding }),
      });

      const chunk: ContentChunk = {
        id: 'test-2',
        path: 'test2.md',
        title: 'Test Title',
        content: 'This is test content',
        tokens: ['test', 'content'],
        firstParagraph: 'First paragraph',
        headings: ['Heading 1', 'Heading 2'],
        frontmatter: { tags: ['test'] },
        wordCount: 3,
        chunkIndex: 0,
        chunkSize: 1,
      };

      const result = await vectorizer.vectorize(chunk);

      expect(result).toEqual(mockEmbedding);
      expect(mockFetch).toHaveBeenCalledWith('http://localhost:11434/api/embeddings', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: expect.stringContaining('"model":"test-model"'),
      });
    });

    it('should handle empty embedding response', async () => {
      mockFetch.mockResolvedValueOnce({
        ok: true,
        json: async () => ({ embedding: [] }),
      });

      const chunk: ContentChunk = {
        id: 'test-3',
        path: 'empty.md',
        title: 'Empty',
        content: 'Content',
        tokens: ['content'],
        firstParagraph: 'First',
        headings: [],
        frontmatter: {},
        wordCount: 1,
        chunkIndex: 0,
        chunkSize: 1,
      };

      const result = await vectorizer.vectorize(chunk);
      expect(result).toEqual([]);
    });

    it('should return empty array on non-200 response', async () => {
      mockFetch.mockResolvedValueOnce({
        ok: false,
        status: 500,
        json: async () => ({}),
      });

      const chunk: ContentChunk = {
        id: 'test-4',
        path: 'error.md',
        title: 'Error',
        content: 'Content',
        tokens: ['content'],
        firstParagraph: 'First',
        headings: [],
        frontmatter: {},
        wordCount: 1,
        chunkIndex: 0,
        chunkSize: 1,
      };

      const result = await vectorizer.vectorize(chunk);
      expect(result).toEqual([]);
    });

    it('should return empty array on invalid JSON response', async () => {
      mockFetch.mockResolvedValueOnce({
        ok: true,
        json: async () => ({ invalid: 'response' }),
      });

      const chunk: ContentChunk = {
        id: 'test-5',
        path: 'invalid.md',
        title: 'Invalid',
        content: 'Content',
        tokens: ['content'],
        firstParagraph: 'First',
        headings: [],
        frontmatter: {},
        wordCount: 1,
        chunkIndex: 0,
        chunkSize: 1,
      };

      const result = await vectorizer.vectorize(chunk);
      expect(result).toEqual([]);
    });

    it('should handle network errors gracefully', async () => {
      mockFetch.mockRejectedValueOnce(new Error('Network error'));

      const chunk: ContentChunk = {
        id: 'test-10',
        path: 'no-frontmatter.md',
        title: 'Test',
        content: 'Content',
        tokens: ['content'],
        firstParagraph: 'First',
        headings: [],
        frontmatter: {},
        wordCount: 1,
        chunkIndex: 0,
        chunkSize: 1,
      };

      const result = await vectorizer.vectorize(chunk);
      expect(result).toEqual([]);
    });
  });

  describe('createPrompt', () => {
    it('should create prompt from all available content', () => {
      // Access the private method through reflection for testing
      const chunk: ContentChunk = {
        id: 'test-7',
        path: 'main.md',
        title: 'Test Title',
        content: 'This is the main content with some text.',
        tokens: ['main', 'content'],
        firstParagraph: 'This is the first paragraph.',
        headings: ['Main Heading', 'Sub Heading'],
        frontmatter: { tags: ['test'], date: '2024-01-01' },
        wordCount: 7,
        chunkIndex: 0,
        chunkSize: 1,
      };

      // Use any to access private method for testing
      const prompt = (vectorizer as any).createPrompt(chunk);
      expect(prompt).toContain('Test Title');
      expect(prompt).toContain('This is the first paragraph');
      expect(prompt).toContain('Main Heading');
      expect(prompt).toContain('Sub Heading');
      // Frontmatter is no longer included in embedding prompts
    });

    it('should handle empty content fields gracefully', () => {
      const chunk: ContentChunk = {
        id: 'test-8',
        path: 'only.md',
        title: '',
        content: 'Only content',
        tokens: ['only', 'content'],
        firstParagraph: '',
        headings: [],
        frontmatter: {},
        wordCount: 1,
        chunkIndex: 0,
        chunkSize: 1,
      };

      const prompt = (vectorizer as any).createPrompt(chunk);
      expect(prompt).toContain('Only content');
      // Frontmatter is no longer included in embedding prompts
    });

    it('should limit content length', () => {
      const longContent = 'a'.repeat(1500);
      const chunk: ContentChunk = {
        id: 'test-9',
        path: 'long.md',
        title: 'Test',
        content: longContent,
        tokens: ['a'],
        firstParagraph: 'First',
        headings: [],
        frontmatter: {},
        wordCount: 1500,
        chunkIndex: 0,
        chunkSize: 1,
      };

      const prompt = (vectorizer as any).createPrompt(chunk);
      expect(prompt).not.toContain('a'.repeat(1500));
      expect(prompt).toContain('a'.repeat(500));
    });

    it('should handle missing frontmatter gracefully', () => {
      const chunk: ContentChunk = {
        id: 'test-11',
        path: 'test.md',
        title: 'Test',
        content: 'Content',
        tokens: ['test'],
        firstParagraph: 'First',
        headings: [],
        frontmatter: {},
        wordCount: 1,
        chunkIndex: 0,
        chunkSize: 1,
      };

      const prompt = (vectorizer as any).createPrompt(chunk);
      expect(prompt).toContain('Test');
      expect(prompt).toContain('Content');
      expect(prompt).toContain('First');
    });
  });

  describe('isEmbeddingResponse', () => {
    it('should validate correct embedding response', () => {
      const response = { embedding: [1, 2, 3] };
      expect((vectorizer as any).isEmbeddingResponse([1, 2, 3])).toBe(false);
    });

    it('should reject non-array embedding', () => {
      const response = { embedding: 'not an array' };
      expect((vectorizer as any).isEmbeddingResponse(response)).toBe(false);
    });

    it('should reject embedding with non-numeric values', () => {
      const response = { embedding: [1, 'two', 3] };
      expect((vectorizer as any).isEmbeddingResponse(response)).toBe(false);
    });

    it('should reject null/undefined', () => {
      expect((vectorizer as any).isEmbeddingResponse(null)).toBe(false);
      expect((vectorizer as any).isEmbeddingResponse(undefined)).toBe(false);
    });

    it('should reject plain array', () => {
      expect((vectorizer as any).isEmbeddingResponse([1, 2, 3])).toBe(false);
    });
  });
});