GitNexus/gitnexus/test/unit/bm25-search.test.ts

import { describe, it, expect, vi, beforeEach } from 'vitest';
import { searchFTSFromLbug, type BM25SearchResult } from '../../src/core/search/bm25-index.js';

vi.mock('../../src/core/lbug/lbug-adapter.js', async (importOriginal) => {
  const actual = await importOriginal<typeof import('../../src/core/lbug/lbug-adapter.js')>();
  return {
    ...actual,
    queryFTS: vi.fn().mockResolvedValue([]),
    createFTSIndex: vi.fn().mockResolvedValue(undefined),
  };
});

// Pool adapter is dynamically imported by the MCP-pool path of
// `searchFTSFromLbug`. We mock it so we can drive the executor without
// spinning up a real LadybugDB pool.
const mockExecuteParameterized = vi.fn();
vi.mock('../../src/core/lbug/pool-adapter.js', () => ({
  executeParameterized: (repoId: string, cypher: string, params: Record<string, any>) =>
    mockExecuteParameterized(repoId, cypher, params),
  addPoolCloseListener: vi.fn(),
}));

describe('BM25 search', () => {
  describe('createSearchFTSIndexes', () => {
    beforeEach(() => {
      vi.clearAllMocks();
    });

    it('creates the configured indexes on the writable analysis path', async () => {
      const { createFTSIndex } = await import('../../src/core/lbug/lbug-adapter.js');
      const { createSearchFTSIndexes } = await import('../../src/core/search/fts-indexes.js');

      await createSearchFTSIndexes();

      expect(vi.mocked(createFTSIndex).mock.calls).toEqual([
        ['File', 'file_fts', ['name', 'content']],
        ['Function', 'function_fts', ['name', 'content']],
        ['Class', 'class_fts', ['name', 'content']],
        ['Method', 'method_fts', ['name', 'content']],
        ['Interface', 'interface_fts', ['name', 'content']],
      ]);
    });

    it('verifies all configured FTS indexes are queryable', async () => {
      const executeQuery = vi.fn().mockResolvedValue([]);
      const { verifySearchFTSIndexes } = await import('../../src/core/search/fts-indexes.js');

      const missing = await verifySearchFTSIndexes(executeQuery);

      expect(missing).toEqual([]);
      expect(executeQuery).toHaveBeenCalledTimes(5);
    });

    it('reports missing indexes when an FTS probe fails', async () => {
      const executeQuery = vi
        .fn()
        .mockResolvedValueOnce([])
        .mockRejectedValueOnce(new Error('index does not exist'))
        .mockResolvedValueOnce([])
        .mockResolvedValueOnce([])
        .mockResolvedValueOnce([]);
      const { verifySearchFTSIndexes } = await import('../../src/core/search/fts-indexes.js');

      const missing = await verifySearchFTSIndexes(executeQuery);

      expect(missing).toEqual(['Function.function_fts']);
    });
  });

  describe('searchFTSFromLbug', () => {
    it('returns empty results when LadybugDB is not initialized', async () => {
      // Simulate an uninitialized DB: queryFTS throws instead of returning rows
      const { queryFTS } = await import('../../src/core/lbug/lbug-adapter.js');
      vi.mocked(queryFTS).mockRejectedValue(new Error('DB not initialized'));

      const { results, ftsAvailable } = await searchFTSFromLbug('test query');
      expect(Array.isArray(results)).toBe(true);
      expect(results).toHaveLength(0);
      expect(ftsAvailable).toBe(false);
    });

    it('handles empty query', async () => {
      const { results } = await searchFTSFromLbug('');
      expect(Array.isArray(results)).toBe(true);
    });

    it('accepts custom limit parameter', async () => {
      const { results } = await searchFTSFromLbug('test', 5);
      expect(Array.isArray(results)).toBe(true);
    });
  });

  describe('BM25SearchResult type', () => {
    it('has correct shape', () => {
      const result: BM25SearchResult = {
        filePath: 'src/index.ts',
        score: 1.5,
        rank: 1,
      };
      expect(result.filePath).toBe('src/index.ts');
      expect(result.score).toBe(1.5);
      expect(result.rank).toBe(1);
    });

    it('accepts optional nodeIds field', () => {
      const result: BM25SearchResult = {
        filePath: 'src/index.ts',
        score: 1.5,
        rank: 1,
        nodeIds: ['func:id1', 'func:id2'],
      };
      expect(result.nodeIds).toEqual(['func:id1', 'func:id2']);
    });
  });

  describe('score aggregation', () => {
    beforeEach(() => {
      vi.clearAllMocks();
    });

    it('sums only top-3 scoring nodes per file when more than 3 match', async () => {
      const { queryFTS } = await import('../../src/core/lbug/lbug-adapter.js');
      // File table: empty; Function table: 5 hits for the same file; rest: empty
      vi.mocked(queryFTS)
        .mockResolvedValueOnce([]) // File
        .mockResolvedValueOnce([
          // Function — 5 hits, scores 10/9/8/7/6
          { filePath: 'src/views.py', score: 10, nodeId: 'func:node1', name: 'get_queryset' },
          { filePath: 'src/views.py', score: 9, nodeId: 'func:node2', name: 'post' },
          { filePath: 'src/views.py', score: 8, nodeId: 'func:node3', name: 'delete' },
          { filePath: 'src/views.py', score: 7, nodeId: 'func:node4', name: 'patch' },
          { filePath: 'src/views.py', score: 6, nodeId: 'func:node5', name: 'put' },
        ])
        .mockResolvedValueOnce([]) // Class
        .mockResolvedValueOnce([]) // Method
        .mockResolvedValueOnce([]); // Interface

      const { results } = await searchFTSFromLbug('queryset');

      expect(results).toHaveLength(1);
      expect(results[0].filePath).toBe('src/views.py');
      // Only top-3 scores (10+9+8=27), not naive sum of all 5 (10+9+8+7+6=40)
      expect(results[0].score).toBe(27);
      expect(results[0].nodeIds).toEqual(['func:node1', 'func:node2', 'func:node3']);
    });

    it('propagates nodeIds for files with fewer than 3 matching nodes', async () => {
      const { queryFTS } = await import('../../src/core/lbug/lbug-adapter.js');
      vi.mocked(queryFTS)
        .mockResolvedValueOnce([]) // File
        .mockResolvedValueOnce([
          // Function — 2 hits
          { filePath: 'src/models.py', score: 5, nodeId: 'func:m1', name: 'save' },
          { filePath: 'src/models.py', score: 3, nodeId: 'func:m2', name: 'delete' },
        ])
        .mockResolvedValueOnce([]) // Class
        .mockResolvedValueOnce([]) // Method
        .mockResolvedValueOnce([]); // Interface

      const { results } = await searchFTSFromLbug('model');

      expect(results).toHaveLength(1);
      expect(results[0].score).toBe(8); // 5+3
      expect(results[0].nodeIds).toEqual(['func:m1', 'func:m2']);
    });

    it('filters out empty nodeIds', async () => {
      const { queryFTS } = await import('../../src/core/lbug/lbug-adapter.js');
      vi.mocked(queryFTS)
        .mockResolvedValueOnce([]) // File
        .mockResolvedValueOnce([
          // Function — nodes with no id
          { filePath: 'src/utils.py', score: 5, nodeId: '', name: 'helper' },
          { filePath: 'src/utils.py', score: 3, nodeId: '', name: 'util' },
        ])
        .mockResolvedValueOnce([]) // Class
        .mockResolvedValueOnce([]) // Method
        .mockResolvedValueOnce([]); // Interface

      const { results } = await searchFTSFromLbug('util');

      expect(results).toHaveLength(1);
      expect(results[0].nodeIds).toEqual([]);
    });

    it('merges hits across multiple index tables for the same file', async () => {
      const { queryFTS } = await import('../../src/core/lbug/lbug-adapter.js');
      vi.mocked(queryFTS)
        .mockResolvedValueOnce([
          // File table
          { filePath: 'src/auth.py', score: 4, nodeId: 'file:auth', name: 'auth.py' },
        ])
        .mockResolvedValueOnce([
          // Function table
          { filePath: 'src/auth.py', score: 9, nodeId: 'func:login', name: 'login' },
        ])
        .mockResolvedValueOnce([
          // Class table
          { filePath: 'src/auth.py', score: 7, nodeId: 'cls:User', name: 'User' },
        ])
        .mockResolvedValueOnce([]) // Method
        .mockResolvedValueOnce([]); // Interface

      const { results } = await searchFTSFromLbug('auth');

      expect(results).toHaveLength(1);
      // All 3 hits (scores 9+7+4=20) — each from a different table, all top-3
      expect(results[0].score).toBe(20);
      expect(results[0].nodeIds).toEqual(['func:login', 'cls:User', 'file:auth']);
    });

    it('ranks files by aggregated score descending', async () => {
      const { queryFTS } = await import('../../src/core/lbug/lbug-adapter.js');
      vi.mocked(queryFTS)
        .mockResolvedValueOnce([]) // File
        .mockResolvedValueOnce([
          // Function — hits across two files
          { filePath: 'src/low.py', score: 2, nodeId: 'func:a', name: 'a' },
          { filePath: 'src/high.py', score: 9, nodeId: 'func:b', name: 'b' },
        ])
        .mockResolvedValueOnce([]) // Class
        .mockResolvedValueOnce([]) // Method
        .mockResolvedValueOnce([]); // Interface

      const { results } = await searchFTSFromLbug('fn');

      expect(results[0].filePath).toBe('src/high.py');
      expect(results[1].filePath).toBe('src/low.py');
      expect(results[0].rank).toBe(1);
      expect(results[1].rank).toBe(2);
    });
  });

  describe('MCP pool path', () => {
    const REPO = 'test-repo-readonly-fts';

    beforeEach(() => {
      mockExecuteParameterized.mockReset();
    });

    it('queries existing FTS indexes without issuing CREATE_FTS_INDEX', async () => {
      mockExecuteParameterized.mockImplementation(
        async (_repo: string, cypher: string, params: Record<string, any>) => {
          if (cypher.includes('CREATE_FTS_INDEX')) {
            throw new Error('query path must stay read-only');
          }

          if (params.query === 'login' && cypher.includes("QUERY_FTS_INDEX('Function'")) {
            return [{ node: { filePath: 'src/auth.ts', id: 'func:login' }, score: 8 }];
          }
          return [];
        },
      );

      const { results } = await searchFTSFromLbug('login', 5, REPO);

      expect(results).toEqual([
        { filePath: 'src/auth.ts', score: 8, rank: 1, nodeIds: ['func:login'] },
      ]);
      expect(
        mockExecuteParameterized.mock.calls.some((c) => String(c[1]).includes('CREATE_FTS_INDEX')),
      ).toBe(false);
    });

    it('binds FTS user query text as a parameter in pool mode', async () => {
      mockExecuteParameterized.mockResolvedValue([]);

      const userQuery = "BrowserWindow create delete set remove 'main' window";
      await searchFTSFromLbug(userQuery, 5, REPO);

      expect(mockExecuteParameterized).toHaveBeenCalled();
      for (const call of mockExecuteParameterized.mock.calls) {
        const cypher = String(call[1]);
        expect(cypher).toContain('$query');
        expect(cypher).not.toContain(userQuery);
        expect(cypher.toUpperCase()).not.toMatch(/\bCREATE\b/);
        expect(cypher.toUpperCase()).not.toMatch(/\bDELETE\b/);
        expect(cypher.toUpperCase()).not.toMatch(/\bSET\b/);
        expect(cypher.toUpperCase()).not.toMatch(/\bREMOVE\b/);
        expect(call[2]).toEqual({ query: userQuery });
      }
    });

    it('uses the configured FTS query set on every call', async () => {
      mockExecuteParameterized.mockResolvedValue([]);

      await searchFTSFromLbug('anything', 5, REPO);

      const queryCalls = mockExecuteParameterized.mock.calls.filter((c) =>
        String(c[1]).includes('QUERY_FTS_INDEX'),
      );
      expect(queryCalls.map((c) => String(c[1]).match(/QUERY_FTS_INDEX\('([^']+)'/)?.[1])).toEqual([
        'File',
        'Function',
        'Class',
        'Method',
        'Interface',
      ]);
    });
  });
});