notebook-ai-extension/src/__tests__/cell_context.spec.ts

import type { CodeCell, ICodeCellModel } from '@jupyterlab/cells';
import type { NotebookPanel } from '@jupyterlab/notebook';

import { collectErrorOutputs, extractCellContext } from '../context/cell_context';

function makeMockOutputs(items: unknown[]): ICodeCellModel['outputs'] {
  return {
    get length() {
      return items.length;
    },
    get(i: number) {
      return items[i];
    },
    // Other IObservableList methods are unused in collectErrorOutputs.
  } as unknown as ICodeCellModel['outputs'];
}

function makeMockModel(overrides: {
  id?: string;
  type?: 'code' | 'markdown' | 'raw';
  source?: string;
  outputs?: unknown[];
}): ICodeCellModel {
  return {
    id: overrides.id ?? 'cell-1',
    type: overrides.type ?? 'code',
    sharedModel: {
      getSource: () => overrides.source ?? '',
    },
    outputs: makeMockOutputs(overrides.outputs ?? []),
  } as unknown as ICodeCellModel;
}

function makeMockCell(source: string, errorOutputs: unknown[] = []): CodeCell {
  return { model: makeMockModel({ source, outputs: errorOutputs }) } as CodeCell;
}

function makeMockNotebook(cells: CodeCell[]): { widgets: CodeCell[] } {
  return { widgets: cells };
}

function makeMockPanel(
  notebook: { widgets: CodeCell[] },
  path: string
): NotebookPanel {
  return {
    context: { path } as NotebookPanel['context'],
    content: notebook as unknown as NotebookPanel['content'],
  } as unknown as NotebookPanel;
}

describe('collectErrorOutputs', () => {
  it('returns null when there are no outputs', () => {
    const model = makeMockModel({ outputs: [] });
    expect(collectErrorOutputs(model)).toBeNull();
  });

  it('returns null when outputs are non-error', () => {
    const model = makeMockModel({
      outputs: [{ type: 'stream', text: 'hello' }],
    });
    expect(collectErrorOutputs(model)).toBeNull();
  });

  it('extracts the first error output', () => {
    const errOut = {
      type: 'error',
      ename: 'ValueError',
      evalue: 'bad value',
      traceback: ['line 1', 'line 2'],
    };
    const model = makeMockModel({ outputs: [errOut] });
    const result = collectErrorOutputs(model);
    expect(result).toEqual({
      ename: 'ValueError',
      evalue: 'bad value',
      traceback: ['line 1', 'line 2'],
    });
  });

  it('handles missing traceback gracefully', () => {
    const model = makeMockModel({
      outputs: [{ type: 'error', ename: 'E', evalue: 'v' }],
    });
    const result = collectErrorOutputs(model);
    expect(result?.traceback).toEqual([]);
  });
});

describe('extractCellContext', () => {
  it('returns null when cell is not in notebook', () => {
    const cell = makeMockCell('x = 1');
    const other = makeMockCell('y = 2');
    const notebook = makeMockNotebook([other]);
    const panel = makeMockPanel(notebook, 'foo.ipynb');
    expect(extractCellContext(cell, panel)).toBeNull();
  });

  it('extracts source, language, and indices', () => {
    const cell = makeMockCell('x = 1');
    const notebook = makeMockNotebook([cell]);
    const panel = makeMockPanel(notebook, 'foo.ipynb');
    const ctx = extractCellContext(cell, panel);
    expect(ctx).toMatchObject({
      notebookPath: 'foo.ipynb',
      cellId: 'cell-1',
      language: 'python',
      cellIndex: 0,
      totalCells: 1,
      source: 'x = 1',
      previousCode: null,
      error: null,
    });
  });

  it('captures the previous cell source as previousCode', () => {
    const prev = makeMockCell('import os');
    const cell = makeMockCell('os.getcwd()');
    const notebook = makeMockNotebook([prev, cell]);
    const panel = makeMockPanel(notebook, 'foo.ipynb');
    const ctx = extractCellContext(cell, panel);
    expect(ctx?.previousCode).toBe('import os');
    expect(ctx?.cellIndex).toBe(1);
    expect(ctx?.totalCells).toBe(2);
  });

  it('includes the error when the cell has an error output', () => {
    const cell = makeMockCell('1/0', [
      {
        type: 'error',
        ename: 'ZeroDivisionError',
        evalue: 'division by zero',
        traceback: ['Traceback...'],
      },
    ]);
    const notebook = makeMockNotebook([cell]);
    const panel = makeMockPanel(notebook, 'foo.ipynb');
    const ctx = extractCellContext(cell, panel);
    expect(ctx?.error).toEqual({
      ename: 'ZeroDivisionError',
      evalue: 'division by zero',
      traceback: ['Traceback...'],
    });
  });

  it('returns null when given null cell or panel', () => {
    expect(
      extractCellContext(
        null as unknown as CodeCell,
        makeMockPanel(makeMockNotebook([]), 'x.ipynb')
      )
    ).toBeNull();
  });
});