GitNexus/gitnexus/test/integration/resolvers/ipynb-python-pipeline.test.ts

/**
 * Jupyter notebooks are indexed as Python by extracting code cells.
 */
import { describe, it, expect } from 'vitest';
import path from 'path';
import fs from 'node:fs';
import os from 'node:os';
import {
  getNodesByLabel,
  getNodesByLabelFull,
  getRelationships,
  runPipelineFromRepo,
  writeFixtureRepo,
} from './helpers.js';
import { extractNotebookPython } from '../../../src/core/ingestion/ipynb-extractor.js';

function pythonNotebook(cells: Array<{ source: string | string[]; language?: string }>): string {
  return JSON.stringify(
    {
      nbformat: 4,
      nbformat_minor: 5,
      metadata: {
        kernelspec: { display_name: 'Python 3', language: 'python', name: 'python3' },
        language_info: { name: 'python' },
      },
      cells: cells.map((c) => ({
        cell_type: 'code',
        metadata: c.language ? { language: c.language } : {},
        source: c.source,
        outputs: [],
      })),
    },
    null,
    2,
  );
}

describe('Jupyter notebook Python pipeline', () => {
  it('indexes notebook functions, imports sibling py, and skip-fails closed', async () => {
    const root = fs.mkdtempSync(path.join(os.tmpdir(), 'gn-ipynb-'));
    try {
      writeFixtureRepo(root, {
        'lib.py': 'def helper():\n    return 1\n',
        'analysis.ipynb': pythonNotebook([
          { source: ['from lib import helper\n'] },
          { source: ['def train():\n', '    return helper()\n'] },
        ]),
        'broken.ipynb': '{not json',
        'julia.ipynb': JSON.stringify({
          nbformat: 4,
          nbformat_minor: 5,
          metadata: {
            kernelspec: { language: 'julia', name: 'julia', display_name: 'Julia' },
          },
          cells: [
            { cell_type: 'code', metadata: {}, source: ['function train() end\n'], outputs: [] },
          ],
        }),
      });
      const result = await runPipelineFromRepo(root, () => {}, { skipGraphPhases: true });
      const functions = getNodesByLabel(result, 'Function');
      expect(functions).toContain('train');
      expect(functions).toContain('helper');
      const train = getNodesByLabelFull(result, 'Function').find((n) => n.name === 'train');
      expect(train?.properties.filePath.replace(/\\/g, '/')).toMatch(/analysis\.ipynb$/);
      const extracted = extractNotebookPython(
        fs.readFileSync(path.join(root, 'analysis.ipynb'), 'utf8'),
      )!;
      expect(train?.properties.startLine).toBe(extracted.segments[1].jsonStartLine);
      const imports = getRelationships(result, 'IMPORTS');
      expect(
        imports.some(
          (e) =>
            e.sourceFilePath.replace(/\\/g, '/').endsWith('analysis.ipynb') &&
            (e.target === 'helper' || e.targetFilePath.replace(/\\/g, '/').endsWith('lib.py')),
        ),
      ).toBe(true);
      expect(
        getNodesByLabelFull(result, 'Function').filter(
          (n) =>
            n.properties.filePath.replace(/\\/g, '/').endsWith('julia.ipynb') && n.name === 'train',
        ),
      ).toHaveLength(0);
      expect(
        getNodesByLabelFull(result, 'Function').filter((n) =>
          n.properties.filePath.replace(/\\/g, '/').endsWith('broken.ipynb'),
        ),
      ).toHaveLength(0);
    } finally {
      fs.rmSync(root, { recursive: true, force: true });
    }
  }, 60000);
});