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

import { ServerConnection } from '@jupyterlab/services';

import { callOpenCodeEdit, callOpenCodeProviders } from '../api/opencode_client';
import type { OpenCodeRequest } from '../types';

jest.mock('@jupyterlab/services', () => {
  const actual = jest.requireActual('@jupyterlab/services');
  return {
    ...actual,
    ServerConnection: {
      ...actual.ServerConnection,
      makeRequest: jest.fn(),
    },
  };
});

const mockedMakeRequest = ServerConnection.makeRequest as jest.MockedFunction<
  typeof ServerConnection.makeRequest
>;

function mockServerSettings(): ServerConnection.ISettings {
  return {
    baseUrl: 'http://localhost:8888',
    wsUrl: 'ws://localhost:8888',
    token: 'test',
    init: { headers: { 'X-Test': '1' } },
    fetch: {} as unknown as typeof fetch,
    RequestClass: {} as unknown as typeof Request,
    Headers: {} as unknown as typeof Headers,
    appendToken: false,
    pageUrl: '',
    settings: {} as unknown as ServerConnection.ISettings,
    displayName: 'JupyterLab',
    handleError: () => undefined,
    requestHeaders: {},
    wsHeaders: {},
    userSettings: {},
  } as unknown as ServerConnection.ISettings;
}

function mockResponse(overrides: {
  ok: boolean;
  status: number;
  statusText: string;
  text: string;
}): Response {
  return {
    ok: overrides.ok,
    status: overrides.status,
    statusText: overrides.statusText,
    text: async () => overrides.text,
    headers: {} as unknown as Headers,
    json: async () => JSON.parse(overrides.text),
    url: '',
    redirected: false,
    type: 'basic',
    body: null,
    bodyUsed: false,
    arrayBuffer: async () => new ArrayBuffer(0),
    blob: async () => new Blob(),
    formData: async () => new FormData(),
    clone: () => mockResponse(overrides),
  } as unknown as Response;
}

const sampleRequest: OpenCodeRequest = {
  prompt: 'add type hints',
  context: {
    notebookPath: 'foo.ipynb',
    cellId: 'cell-1',
    language: 'python',
    cellIndex: 0,
    totalCells: 1,
    source: 'def foo(x): return x',
    previousCode: null,
    error: null,
  },
};

describe('callOpenCodeEdit', () => {
  beforeEach(() => {
    mockedMakeRequest.mockReset();
  });

  it('POSTs to /opencode-bridge/edit with JSON body', async () => {
    const respBody = {
      ok: true,
      markdown: '```python\ndef foo(x: int) -> int: return x\n```',
      sessionId: 'sid',
      notebookPath: 'foo.ipynb',
    };
    mockedMakeRequest.mockResolvedValue(
      mockResponse({ ok: true, status: 200, statusText: 'OK', text: JSON.stringify(respBody) })
    );

    const settings = mockServerSettings();
    const resp = await callOpenCodeEdit(sampleRequest, settings);

    expect(mockedMakeRequest).toHaveBeenCalledWith(
      'http://localhost:8888/opencode-bridge/edit',
      expect.objectContaining({
        method: 'POST',
        body: JSON.stringify(sampleRequest),
        headers: { 'Content-Type': 'application/json' },
      }),
      settings
    );
    expect(resp.ok).toBe(true);
    if (resp.ok) {
      expect(resp.markdown).toContain('int');
    }
  });

  it('returns failure response when server returns ok: false', async () => {
    mockedMakeRequest.mockResolvedValue(
      mockResponse({
        ok: false,
        status: 502,
        statusText: 'Bad Gateway',
        text: JSON.stringify({ ok: false, error: 'opencode down' }),
      })
    );

    const resp = await callOpenCodeEdit(sampleRequest, mockServerSettings());
    expect(resp.ok).toBe(false);
    if (!resp.ok) {
      expect(resp.error).toBe('opencode down');
    }
  });

  it('throws on network error', async () => {
    mockedMakeRequest.mockRejectedValue(new Error('ECONNREFUSED'));
    await expect(
      callOpenCodeEdit(sampleRequest, mockServerSettings())
    ).rejects.toThrow(/network error/);
  });

  it('throws on empty response body', async () => {
    mockedMakeRequest.mockResolvedValue(
      mockResponse({ ok: true, status: 200, statusText: 'OK', text: '' })
    );
    await expect(
      callOpenCodeEdit(sampleRequest, mockServerSettings())
    ).rejects.toThrow(/empty response/);
  });
});

describe('callOpenCodeProviders', () => {
  beforeEach(() => {
    mockedMakeRequest.mockReset();
  });

  it('GETs /opencode-bridge/providers and returns parsed JSON', async () => {
    const respBody = {
      providers: [
        { id: 'anthropic', models: [{ id: 'claude-sonnet-4-20250514' }] },
        { id: 'openai', models: [{ id: 'gpt-4' }, { id: 'gpt-3.5-turbo' }] }
      ]
    };
    mockedMakeRequest.mockResolvedValue({
      ok: true,
      status: 200,
      statusText: 'OK',
      text: async () => JSON.stringify(respBody),
      json: async () => respBody,
      headers: {} as any,
      url: '',
      redirected: false,
      type: 'basic',
      body: null,
      bodyUsed: false,
      arrayBuffer: async () => new ArrayBuffer(0),
      blob: async () => new Blob(),
      formData: async () => new FormData(),
      clone: function () { return this; }
    } as any);

    const settings = mockServerSettings();
    const resp = await callOpenCodeProviders(settings);

    expect(mockedMakeRequest).toHaveBeenCalledWith(
      'http://localhost:8888/opencode-bridge/providers',
      expect.objectContaining({ method: 'GET' }),
      settings
    );
    expect(resp.providers).toHaveLength(2);
    expect(resp.providers[0].id).toBe('anthropic');
    expect(resp.providers[1].models).toHaveLength(2);
  });

  it('throws on non-2xx response', async () => {
    mockedMakeRequest.mockResolvedValue({
      ok: false,
      status: 502,
      statusText: 'Bad Gateway',
      text: async () => '',
      json: async () => null,
      headers: {} as any,
      url: '',
      redirected: false,
      type: 'basic',
      body: null,
      bodyUsed: false,
      arrayBuffer: async () => new ArrayBuffer(0),
      blob: async () => new Blob(),
      formData: async () => new FormData(),
      clone: function () { return this; }
    } as any);

    await expect(
      callOpenCodeProviders(mockServerSettings())
    ).rejects.toThrow(/502/);
  });

  it('throws on network error', async () => {
    mockedMakeRequest.mockRejectedValue(new Error('ECONNREFUSED'));
    await expect(
      callOpenCodeProviders(mockServerSettings())
    ).rejects.toThrow(/network error/);
  });
});