GitNexus/gitnexus/test/unit/ast-utils.test.ts

import { beforeEach, describe, expect, it, vi } from 'vitest';

const { createParserForLanguage, getLanguageFromFilename, parseSourceSafeSpy } = vi.hoisted(() => ({
  createParserForLanguage: vi.fn(),
  getLanguageFromFilename: vi.fn((filePath: string) =>
    filePath.endsWith('.py') ? 'python' : 'typescript',
  ),
  parseSourceSafeSpy: vi.fn(),
}));

vi.mock('../../src/core/tree-sitter/parser-loader.js', () => ({
  createParserForLanguage,
  isLanguageAvailable: vi.fn().mockReturnValue(true),
  resolveLanguageKey: vi.fn((language: string, filePath?: string) =>
    language === 'typescript' && filePath?.endsWith('.tsx') ? 'typescript:tsx' : language,
  ),
}));

vi.mock('../../src/core/tree-sitter/safe-parse.js', async () => {
  const { buildSafeParseMock } = await import('../helpers/parse-source-safe-mock.js');
  return buildSafeParseMock(parseSourceSafeSpy);
});

// Partial mock: `ast-utils` now resolves the LanguageProvider registry to apply
// `preprocessSource`, and that graph needs the real shared exports (#2771).
vi.mock('gitnexus-shared', async (importOriginal) => ({
  ...(await importOriginal<typeof import('gitnexus-shared')>()),
  getLanguageFromFilename,
}));

describe('ensureAndParse', () => {
  beforeEach(() => {
    vi.resetModules();
    createParserForLanguage.mockReset();
    getLanguageFromFilename.mockClear();
  });

  it('reuses the parser for the same grammar key across interleaved languages', async () => {
    const tsParse = vi
      .fn()
      .mockReturnValueOnce({ lang: 'ts', content: 'first' })
      .mockReturnValueOnce({ lang: 'ts', content: 'second' });
    const pyParse = vi.fn().mockReturnValue({ lang: 'py', content: 'middle' });

    createParserForLanguage.mockImplementation(async (language: string, filePath?: string) => {
      if (language === 'typescript') return { parse: tsParse, key: filePath };
      if (language === 'python') return { parse: pyParse, key: filePath };
      throw new Error(`unexpected language ${language}`);
    });

    const { ensureAndParse } = await import('../../src/core/embeddings/ast-utils.js');

    const tsFirst = await ensureAndParse('const one = 1;', 'first.ts');
    const pyMiddle = await ensureAndParse('value = 1', 'middle.py');
    const tsSecond = await ensureAndParse('const two = 2;', 'second.ts');

    expect(tsFirst).toEqual({ lang: 'ts', content: 'first' });
    expect(pyMiddle).toEqual({ lang: 'py', content: 'middle' });
    expect(tsSecond).toEqual({ lang: 'ts', content: 'second' });
    expect(createParserForLanguage).toHaveBeenCalledTimes(2);
    expect(tsParse).toHaveBeenCalledTimes(2);
    expect(pyParse).toHaveBeenCalledTimes(1);
  });

  it('uses separate parser instances for .ts and .tsx', async () => {
    const tsParse = vi.fn().mockReturnValue({ lang: 'ts' });
    const tsxParse = vi.fn().mockReturnValue({ lang: 'tsx' });

    createParserForLanguage.mockImplementation(async (_language: string, filePath?: string) => {
      if (filePath?.endsWith('.tsx')) return { parse: tsxParse };
      return { parse: tsParse };
    });

    const { ensureAndParse } = await import('../../src/core/embeddings/ast-utils.js');

    await ensureAndParse('const value = 1;', 'plain.ts');
    await ensureAndParse('export const View = <div />;', 'view.tsx');
    await ensureAndParse('const other = 2;', 'other.ts');

    expect(createParserForLanguage).toHaveBeenCalledTimes(2);
    expect(tsParse).toHaveBeenCalledTimes(2);
    expect(tsxParse).toHaveBeenCalledTimes(1);
  });

  // Windows SIGSEGV regression: ensureAndParse must route through parseSourceSafe
  // so >32 767-char inputs do not crash the process. Direct parser.parse(content)
  // on strings that size SIGSEGVs on Windows; the spy assertion is what catches
  // a bypass since parser.parse(40 000 chars) succeeds on Linux/macOS.
  it('routes >32 767-char input through parseSourceSafe', async () => {
    parseSourceSafeSpy.mockClear();

    const fakeParse = vi.fn().mockReturnValue({ rootNode: { type: 'module' } });
    createParserForLanguage.mockResolvedValue({ parse: fakeParse });

    const { ensureAndParse } = await import('../../src/core/embeddings/ast-utils.js');

    const largeInput = 'const x = 1;\n'.repeat(4000); // ~52 000 chars
    expect(largeInput.length).toBeGreaterThan(40_000);

    const result = await ensureAndParse(largeInput, 'big.ts');

    expect(parseSourceSafeSpy).toHaveBeenCalled();
    expect(result).not.toBeNull();
  });

  it('parses Objective-C .h declarations and method snippets with the objc grammar', async () => {
    const objcParse = vi.fn().mockReturnValue({ lang: 'objc' });
    const cppParse = vi.fn().mockReturnValue({ lang: 'cpp' });
    createParserForLanguage.mockImplementation(async (language: string) => {
      if (language === 'objective-c') return { parse: objcParse };
      if (language === 'cpp') return { parse: cppParse };
      throw new Error(`unexpected language ${language}`);
    });

    const { ensureAndParse } = await import('../../src/core/embeddings/ast-utils.js');

    await ensureAndParse('@interface Worker\n- (void)run;\n@end\n', 'Worker.h');
    await ensureAndParse('- (void)run;\n', 'Worker.h');
    await ensureAndParse('class Widget { int value; };\n', 'widget.h');

    expect(createParserForLanguage).toHaveBeenCalledWith('objective-c', 'Worker.h');
    expect(createParserForLanguage).toHaveBeenCalledWith('cpp', 'widget.h');
    expect(objcParse).toHaveBeenCalledTimes(2);
    expect(cppParse).toHaveBeenCalledTimes(1);
  });

  it('parses extracted notebook Python and returns null when extraction fails', async () => {
    parseSourceSafeSpy.mockClear();
    const pyParse = vi.fn().mockReturnValue({ lang: 'py', rootNode: { type: 'module' } });
    createParserForLanguage.mockImplementation(async (language: string) => {
      if (language === 'python') return { parse: pyParse };
      throw new Error(`unexpected language ${language}`);
    });

    const { ensureAndParse } = await import('../../src/core/embeddings/ast-utils.js');
    const nb = JSON.stringify({
      nbformat: 4,
      nbformat_minor: 5,
      metadata: { kernelspec: { language: 'python', name: 'python3', display_name: 'Python' } },
      cells: [
        { cell_type: 'code', metadata: {}, source: ['def train():\n', '    pass\n'], outputs: [] },
      ],
    });

    await ensureAndParse(nb, 'analysis.ipynb');
    expect(parseSourceSafeSpy).toHaveBeenCalled();
    const parsedText = parseSourceSafeSpy.mock.calls.at(-1)?.[1] as string;
    expect(parsedText).toContain('def train');
    expect(parsedText).not.toContain('cell_type');

    parseSourceSafeSpy.mockClear();
    expect(await ensureAndParse('{not json', 'broken.ipynb')).toBeNull();
    expect(parseSourceSafeSpy).not.toHaveBeenCalled();
  });
});