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

/**
 * Unit tests for OpenCodeCellActions (single AI button) and OpenCodeInlinePrompt
 * (inline input panel inside the cell). v3-final: no settings — the model
 * picker is fully dynamic; default selection is the first provider/model.
 *
 * Mocks @jupyterlab/cells, @jupyterlab/apputils, and @lumino/widgets at the
 * test boundary so the real ESM packages are not loaded.
 */
jest.mock('@jupyterlab/cells', () => ({
  CodeCell: class CodeCell {}
}));

jest.mock('@jupyterlab/apputils', () => ({
  Notification: {
    info: jest.fn(),
    error: jest.fn(),
    success: jest.fn(),
    warning: jest.fn()
  }
}));

jest.mock('@lumino/widgets', () => {
  class Widget {
    public node: HTMLElement = document.createElement('div');
    public id = '';
    public parent: Widget | null = null;
    private _isDisposed = false;
    public get isDisposed(): boolean {
      return this._isDisposed;
    }
    public addClass(cls: string): void {
      this.node.classList.add(cls);
    }
    public dispose(): void {
      this._isDisposed = true;
      if (this.node.parentNode) {
        this.node.parentNode.removeChild(this.node);
      }
    }
    public static attach(widget: Widget, host: HTMLElement): void {
      host.appendChild(widget.node);
    }
    public onAfterAttach(_msg: any): void {
      // no-op
    }
  }
  return { Widget };
});

jest.mock('marked', () => ({
  __esModule: true,
  marked: {
    // Minimal but realistic mock: produce a real <pre><code> for
    // fenced code blocks (so the per-code-block toolbar
    // post-processing can find them) and <h1>/<p> for headings/text.
    // Not a full markdown renderer — just enough for the prompt tests.
    parse: (md: string) => {
      let html = '';
      let inFence = false;
      let fenceLang = '';
      let fenceBuf: string[] = [];
      const lines = md.split('\n');
      for (const line of lines) {
        const open = !inFence && /^```(\w*)\s*$/.exec(line);
        const close = inFence && /^```\s*$/.test(line);
        if (open) {
          inFence = true;
          fenceLang = open[1] || '';
          fenceBuf = [];
          continue;
        }
        if (close) {
          html +=
            '<pre><code class="language-' +
            fenceLang +
            '">' +
            fenceBuf.join('\n') +
            '</code></pre>';
          inFence = false;
          continue;
        }
        if (inFence) {
          fenceBuf.push(line);
          continue;
        }
        if (/^# /.test(line)) {
          html += '<h1>' + line.slice(2) + '</h1>';
        } else if (line.trim()) {
          html += '<p>' + line + '</p>';
        }
      }
      return html || '<p>' + md + '</p>';
    }
  }
}));

// jsdom doesn't implement navigator.clipboard; stub it so the copy
// button's happy path can be exercised in tests.
beforeAll(() => {
  Object.defineProperty(globalThis.navigator, 'clipboard', {
    value: { writeText: jest.fn().mockResolvedValue(undefined) },
    configurable: true
  });
});

// Mock the network-touching API module so jsdom tests don't try to hit
// a real notebook server (and to keep the refresh-history logic decoupled
// from the network). The cell_actions tests focus on widget behavior.
jest.mock('../api/opencode_client', () => ({
  callOpenCodeEdit: jest.fn(),
  callOpenCodeProviders: jest.fn(),
  callOpenCodeSessionMessages: jest.fn().mockResolvedValue({ messages: [] })
}));

import {
  OpenCodeCellActions,
  setOpenCodeProviders,
  setOpenCodeServerSettings
} from '../components/opencode_cell_actions';
import { OpenCodeInlinePrompt } from '../components/opencode_inline_prompt';
import { CodeCell } from '@jupyterlab/cells';

function makeFakeOutputs(items: any[]): any {
  return {
    get length() {
      return items.length;
    },
    get(i: number) {
      return items[i];
    }
  };
}

function makeFakeCell(
  source: string,
  errorOutputs: any[] = [],
  notebookPath = 'foo.ipynb',
  cellIndex = 0,
  totalCells = 1
): CodeCell {
  const model: any = {
    id: 'cell-test',
    type: 'code',
    sharedModel: {
      getSource: () => source,
      setSource: jest.fn()
    },
    outputs: makeFakeOutputs(errorOutputs),
    contentChanged: { connect: jest.fn(), disconnect: jest.fn() },
    stateChanged: { connect: jest.fn(), disconnect: jest.fn() }
  };
  const cell = new (CodeCell as unknown as new () => CodeCell)();
  (cell as any).model = model;
  (cell as any).node = document.createElement('div');
  // Minimal editor stub: cursor always at (line 0, column 0) so the
  // "insert at cursor" tests produce a deterministic offset.
  (cell as any).editor = {
    getCursorPosition: () => ({ line: 0, column: 0 })
  };

  const notebook: any = { widgets: [] as CodeCell[] };
  const panel: any = {
    context: { path: notebookPath },
    content: notebook
  };
  for (let i = 0; i < cellIndex; i++) {
    notebook.widgets.push({
      model: { sharedModel: { getSource: () => '' } }
    } as any);
  }
  notebook.widgets.push(cell);
  for (let i = cellIndex + 1; i < totalCells; i++) {
    notebook.widgets.push({
      model: { sharedModel: { getSource: () => '' } }
    } as any);
  }

  Object.defineProperty(cell, 'parent', {
    value: notebook,
    configurable: true
  });
  Object.defineProperty(notebook, 'parent', {
    value: panel,
    configurable: true
  });

  return cell as CodeCell;
}

describe('OpenCodeCellActions', () => {
  it('allows server settings to be injected', () => {
    setOpenCodeServerSettings({} as any);
  });

  it('renders a single AI button', () => {
    const cell = makeFakeCell('x = 1');
    const actions = new OpenCodeCellActions(cell);
    const btns = actions.node.querySelectorAll('button');
    expect(btns.length).toBe(1);
    expect(btns[0].textContent).toContain('AI');
  });

  it('disables the button when notebook panel cannot be resolved', () => {
    const cell = new (CodeCell as unknown as new () => CodeCell)();
    (cell as any).model = {
      id: 'orphan',
      type: 'code',
      sharedModel: { getSource: () => 'x = 1', setSource: jest.fn() },
      outputs: makeFakeOutputs([]),
      contentChanged: { connect: jest.fn(), disconnect: jest.fn() },
      stateChanged: { connect: jest.fn(), disconnect: jest.fn() }
    };
    const actions = new OpenCodeCellActions(cell);
    const btn = actions.node.querySelector('button') as HTMLButtonElement;
    expect(btn.disabled).toBe(true);
  });

  it('clicking the button attaches an OpenCodeInlinePrompt to the cell', () => {
    const cell = makeFakeCell('x = 1');
    const actions = new OpenCodeCellActions(cell);
    Object.defineProperty(actions, 'parent', { value: cell, configurable: true });
    (actions as any).onAfterAttach({} as any);

    const btn = actions.node.querySelector('button') as HTMLButtonElement;
    btn.click();

    const prompt = cell.node.querySelector(
      '.opencode-inline-prompt'
    ) as HTMLElement;
    expect(prompt).not.toBeNull();
  });

  it('clicking the button twice detaches the inline prompt', () => {
    const cell = makeFakeCell('x = 1');
    const actions = new OpenCodeCellActions(cell);
    Object.defineProperty(actions, 'parent', { value: cell, configurable: true });
    (actions as any).onAfterAttach({} as any);

    const btn = actions.node.querySelector('button') as HTMLButtonElement;
    btn.click();
    btn.click();
    expect(cell.node.querySelector('.opencode-inline-prompt')).toBeNull();
  });

  it('disconnects model signals on dispose', () => {
    const cell = makeFakeCell('x = 1');
    const actions = new OpenCodeCellActions(cell);
    actions.dispose();
    const model = (cell as any).model;
    expect(model.contentChanged.disconnect).toHaveBeenCalled();
    expect(model.stateChanged.disconnect).toHaveBeenCalled();
  });

  it('renders provider and model selects; model default uses default[provider] when it matches', () => {
    setOpenCodeProviders({
      providers: [
        {
          id: 'anthropic',
          name: 'Anthropic',
          models: {
            'claude-sonnet-4-20250514': { id: 'claude-sonnet-4-20250514' },
            'claude-haiku-4-5': { id: 'claude-haiku-4-5' }
          }
        },
        { id: 'openai', models: { 'gpt-5': { id: 'gpt-5' } } }
      ],
      default: {
        anthropic: 'claude-sonnet-4-20250514',
        openai: 'gpt-5'
      }
    });
    const cell = makeFakeCell('x = 1');
    const actions = new OpenCodeCellActions(cell);
    Object.defineProperty(actions, 'parent', { value: cell, configurable: true });
    (actions as any).onAfterAttach({} as any);

    const btn = actions.node.querySelector('button') as HTMLButtonElement;
    btn.click();

    const prompt = cell.node.querySelector(
      '.opencode-inline-prompt'
    ) as HTMLElement;
    const providerSel = prompt.querySelector(
      '.opencode-provider-select'
    ) as HTMLSelectElement;
    const modelSel = prompt.querySelector(
      '.opencode-model-select'
    ) as HTMLSelectElement;
    expect(providerSel).not.toBeNull();
    expect(modelSel).not.toBeNull();

    // Provider: first provider selected by default.
    expect(providerSel.options.length).toBe(2);
    expect(providerSel.options[0].value).toBe('anthropic');
    expect(providerSel.options[0].textContent).toContain('anthropic');
    expect(providerSel.options[0].textContent).toContain('Anthropic');
    expect(providerSel.value).toBe('anthropic');

    // Model: options are that provider's model keys, and the default
    // matches default['anthropic'] = 'claude-sonnet-4-20250514'.
    expect(modelSel.options.length).toBe(2);
    expect(modelSel.options[0].value).toBe('claude-sonnet-4-20250514');
    expect(modelSel.options[1].value).toBe('claude-haiku-4-5');
    expect(modelSel.value).toBe('claude-sonnet-4-20250514');

    setOpenCodeProviders(null);
  });

  it('falls back to the first model when default[provider] is not in models', () => {
    setOpenCodeProviders({
      providers: [
        { id: 'bailian', models: { 'kimi/kimi-k2.7-code': { id: 'kimi/kimi-k2.7-code' } } }
      ],
      // default['bailian'] references a modelID not present in models.
      default: { bailian: 'qwen3.7-plus' }
    });
    const cell = makeFakeCell('x = 1');
    const actions = new OpenCodeCellActions(cell);
    Object.defineProperty(actions, 'parent', { value: cell, configurable: true });
    (actions as any).onAfterAttach({} as any);

    const btn = actions.node.querySelector('button') as HTMLButtonElement;
    btn.click();

    const modelSel = cell.node.querySelector(
      '.opencode-inline-prompt .opencode-model-select'
    ) as HTMLSelectElement;
    expect(modelSel.value).toBe('kimi/kimi-k2.7-code');
    setOpenCodeProviders(null);
  });

  it('rebuilds the model select when the provider select changes', () => {
    setOpenCodeProviders({
      providers: [
        {
          id: 'anthropic',
          models: { 'claude-sonnet-4-20250514': { id: 'claude-sonnet-4-20250514' } }
        },
        {
          id: 'openai',
          models: { 'gpt-5': { id: 'gpt-5' } }
        }
      ],
      default: { openai: 'gpt-5' }
    });
    const cell = makeFakeCell('x = 1');
    const actions = new OpenCodeCellActions(cell);
    Object.defineProperty(actions, 'parent', { value: cell, configurable: true });
    (actions as any).onAfterAttach({} as any);

    const btn = actions.node.querySelector('button') as HTMLButtonElement;
    btn.click();

    const prompt = cell.node.querySelector(
      '.opencode-inline-prompt'
    ) as HTMLElement;
    const providerSel = prompt.querySelector(
      '.opencode-provider-select'
    ) as HTMLSelectElement;
    const modelSel = prompt.querySelector(
      '.opencode-model-select'
    ) as HTMLSelectElement;

    // Initial: anthropic, its only model.
    expect(providerSel.value).toBe('anthropic');
    expect(modelSel.options.length).toBe(1);
    expect(modelSel.value).toBe('claude-sonnet-4-20250514');

    // Switch to openai: model select rebuilds with openai's model and
    // default['openai'] = 'gpt-5' is in openai's models so it matches.
    providerSel.value = 'openai';
    providerSel.dispatchEvent(new Event('change'));
    expect(modelSel.options.length).toBe(1);
    expect(modelSel.options[0].value).toBe('gpt-5');
    expect(modelSel.value).toBe('gpt-5');

    setOpenCodeProviders(null);
  });

  it('does not render the selector when providers are not cached', () => {
    setOpenCodeProviders(null);
    const cell = makeFakeCell('x = 1');
    const actions = new OpenCodeCellActions(cell);
    Object.defineProperty(actions, 'parent', { value: cell, configurable: true });
    (actions as any).onAfterAttach({} as any);

    const btn = actions.node.querySelector('button') as HTMLButtonElement;
    btn.click();

    expect(
      cell.node.querySelector('.opencode-inline-prompt .opencode-provider-select')
    ).toBeNull();
    expect(
      cell.node.querySelector('.opencode-inline-prompt .opencode-model-select')
    ).toBeNull();
  });

  it('on a successful response, does NOT replace the cell source and triggers a history refresh', () => {
    setOpenCodeProviders(null);
    const cell = makeFakeCell('x = 1');
    const actions = new OpenCodeCellActions(cell);
    // Click once to create the inline prompt (so the history refresh
    // has a target and the cell source check has a baseline).
    Object.defineProperty(actions, 'parent', { value: cell, configurable: true });
    (actions as any).onAfterAttach({} as any);
    const btn = actions.node.querySelector('button') as HTMLButtonElement;
    btn.click();

    // The history refresh on show is async; spy on _refreshHistory to
    // confirm it gets called (without depending on the network).
    const refreshSpy = jest.spyOn(actions as any, '_refreshHistory');

    (actions as any)._handleResponse({
      ok: true,
      markdown: '# AI says hi\n\n```python\nprint("hi")\n```',
      sessionId: 'ses-abc',
      notebookPath: 'foo.ipynb'
    });

    // The cell source must remain unchanged.
    expect((cell as any).model.sharedModel.setSource).not.toHaveBeenCalled();
    // A history refresh was triggered (to pick up the new assistant msg).
    expect(refreshSpy).toHaveBeenCalledWith('foo.ipynb');
  });

  it('clears the input textarea after a successful response', () => {
    setOpenCodeProviders(null);
    const cell = makeFakeCell('x = 1');
    const actions = new OpenCodeCellActions(cell);
    Object.defineProperty(actions, 'parent', { value: cell, configurable: true });
    (actions as any).onAfterAttach({} as any);

    // Open the prompt and type something into the textarea.
    const aiBtn = actions.node.querySelector('button') as HTMLButtonElement;
    aiBtn.click();
    const prompt = (actions as any)._prompt;
    expect(prompt).not.toBeNull();
    const textarea = prompt._textarea as HTMLTextAreaElement;
    textarea.value = 'make it faster';

    // A successful response clears the input so the user can type a
    // follow-up message without manually deleting the previous prompt.
    (actions as any)._handleResponse({
      ok: true,
      markdown: 'ok',
      sessionId: 'ses-1',
      notebookPath: 'foo.ipynb'
    });
    expect(textarea.value).toBe('');
  });
});

describe('OpenCodeInlinePrompt', () => {
  it('renders a textarea, a send and a cancel button (no selector when providers null)', () => {
    const cell = makeFakeCell('x = 1');
    const prompt = new OpenCodeInlinePrompt(cell, {
      onSubmit: jest.fn(),
      onCancel: jest.fn(),
      disabled: false,
      providers: null
    });
    expect(prompt.node.querySelector('textarea')).not.toBeNull();
    // 2 buttons: 发送 and 取消 (no more output-area close button; the
    // scrollable history area is the display).
    expect(prompt.node.querySelectorAll('button').length).toBe(2);
  });

  it('renders a history area; setMessages renders user text + assistant markdown and scrolls to bottom', () => {
    const cell = makeFakeCell('x = 1');
    const prompt = new OpenCodeInlinePrompt(cell, {
      onSubmit: jest.fn(),
      onCancel: jest.fn(),
      disabled: false,
      providers: null
    });
    const history = prompt.node.querySelector(
      '.opencode-inline-prompt .opencode-inline-history'
    ) as HTMLElement;
    expect(history).not.toBeNull();
    // Initially empty (no messages yet).
    expect(history.children.length).toBe(0);

    prompt.setMessages([
      { role: 'user', content: 'fix the bug' },
      {
        role: 'assistant',
        content: '# Here you go\n\n```python\nx = 1\n```'
      }
    ]);
    expect(history.children.length).toBe(2);
    const userMsg = history.querySelector('.opencode-msg-user') as HTMLElement;
    const asstMsg = history.querySelector('.opencode-msg-assistant') as HTMLElement;
    expect(userMsg).not.toBeNull();
    // User messages are plain text (not rendered as HTML).
    expect(userMsg.textContent).toBe('fix the bug');
    expect(userMsg.innerHTML).not.toContain('<');
    // Assistant content: the marked mock produces a real <h1> + <pre>.
    expect(asstMsg).not.toBeNull();
    expect(asstMsg.innerHTML).toContain('<h1>Here you go</h1>');
    expect(asstMsg.innerHTML).toContain('<pre>');
  });

  it('renders Copy/Insert/Replace buttons INSIDE each code block (not above the message)', () => {
    const cell = makeFakeCell('x = 1');
    const prompt = new OpenCodeInlinePrompt(cell, {
      onSubmit: jest.fn(),
      onCancel: jest.fn(),
      disabled: false,
      providers: null
    });
    // User message: no toolbar.
    prompt.setMessages([{ role: 'user', content: 'help' }]);
    expect(
      prompt.node.querySelector('.opencode-msg-user .opencode-code-toolbar')
    ).toBeNull();

    // Assistant with a fenced code block: toolbar inside .opencode-code-block.
    prompt.setMessages([
      { role: 'assistant', content: '```py\nprint(1)\n```' }
    ]);
    const asst = prompt.node.querySelector(
      '.opencode-msg-assistant'
    ) as HTMLElement;
    // The message-level toolbar must NOT exist.
    expect(asst.querySelector('.opencode-msg-toolbar')).toBeNull();
    // The per-code-block wrapper + toolbar.
    const block = asst.querySelector('.opencode-code-block') as HTMLElement;
    expect(block).not.toBeNull();
    const toolbar = block.querySelector('.opencode-code-toolbar') as HTMLElement;
    expect(toolbar).not.toBeNull();
    const buttons = toolbar.querySelectorAll('button');
    expect(buttons.length).toBe(3);
    expect(buttons[0].textContent).toContain('复制');
    expect(buttons[1].textContent).toContain('插入');
    expect(buttons[2].textContent).toContain('替换');
    // The wrapped <pre> is inside the block, below the toolbar.
    expect(block.querySelector('pre')).not.toBeNull();
  });

  it('assistant with NO code fences has no toolbar', () => {
    const cell = makeFakeCell('x = 1');
    const prompt = new OpenCodeInlinePrompt(cell, {
      onSubmit: jest.fn(),
      onCancel: jest.fn(),
      disabled: false,
      providers: null
    });
    prompt.setMessages([
      { role: 'assistant', content: 'no code here, just an explanation' }
    ]);
    const asst = prompt.node.querySelector(
      '.opencode-msg-assistant'
    ) as HTMLElement;
    expect(asst.querySelector('.opencode-code-toolbar')).toBeNull();
    expect(asst.querySelector('.opencode-code-block')).toBeNull();
  });

  it('Replace button overwrites the cell source with the code block content', () => {
    const cell = makeFakeCell('x = 1');
    const prompt = new OpenCodeInlinePrompt(cell, {
      onSubmit: jest.fn(),
      onCancel: jest.fn(),
      disabled: false,
      providers: null
    });
    prompt.setMessages([
      { role: 'assistant', content: '```py\nprint(1)\n```' }
    ]);
    const replaceBtn = prompt.node.querySelector(
      '.opencode-msg-assistant .opencode-code-block .opencode-code-toolbar button:nth-child(3)'
    ) as HTMLButtonElement;
    replaceBtn.click();

    const setSource = (cell as any).model.sharedModel.setSource as jest.Mock;
    expect(setSource).toHaveBeenCalledTimes(1);
    // The button acts on the code block's text, not the whole markdown
    // (no surrounding ```fences```).
    expect(setSource).toHaveBeenCalledWith('print(1)');
  });

  it('Insert button splices the code block at the editor cursor (offset 0)', () => {
    const cell = makeFakeCell('x = 1');
    const prompt = new OpenCodeInlinePrompt(cell, {
      onSubmit: jest.fn(),
      onCancel: jest.fn(),
      disabled: false,
      providers: null
    });
    prompt.setMessages([
      { role: 'assistant', content: '```py\nprint(1)\n```' }
    ]);
    const insertBtn = prompt.node.querySelector(
      '.opencode-msg-assistant .opencode-code-block .opencode-code-toolbar button:nth-child(2)'
    ) as HTMLButtonElement;
    insertBtn.click();

    const setSource = (cell as any).model.sharedModel.setSource as jest.Mock;
    expect(setSource).toHaveBeenCalledTimes(1);
    // cursor at (0,0) with source "x = 1" -> code "print(1)" inserted at 0
    expect(setSource).toHaveBeenCalledWith('print(1)x = 1');
  });

  it('Copy button writes the code block content to the clipboard', () => {
    const cell = makeFakeCell('x = 1');
    const prompt = new OpenCodeInlinePrompt(cell, {
      onSubmit: jest.fn(),
      onCancel: jest.fn(),
      disabled: false,
      providers: null
    });
    prompt.setMessages([
      { role: 'assistant', content: '```py\nprint(1)\n```' }
    ]);
    const copyBtn = prompt.node.querySelector(
      '.opencode-msg-assistant .opencode-code-block .opencode-code-toolbar button:nth-child(1)'
    ) as HTMLButtonElement;
    copyBtn.click();

    const writeText = (navigator.clipboard.writeText as unknown as jest.Mock);
    expect(writeText).toHaveBeenCalledWith('print(1)');
  });
});