code-review/app/llm.test.js

import { describe, it, beforeEach, afterEach, mock } from 'node:test';
import assert from 'node:assert/strict';

// Mock axios before importing llm.js
import axios from 'axios';

const ENV_KEYS = [
  'OPENAI_API_KEY', 'OPENAI_BASE_URL', 'OPENAI_MODEL',
  'GEMINI_API_KEY', 'GEMINI_BASE_URL', 'GEMINI_MODEL',
  'CLAUDE_API_KEY', 'CLAUDE_BASE_URL', 'CLAUDE_MODEL',
  'OLLAMA_BASE_URL', 'OLLAMA_MODEL',
  'AMAZONQ_API_KEY', 'AMAZONQ_BASE_URL', 'AMAZONQ_MODEL',
];

let saved = {};
beforeEach(() => {
  saved = {};
  for (const k of ENV_KEYS) { saved[k] = process.env[k]; delete process.env[k]; }
});
afterEach(() => {
  for (const k of ENV_KEYS) {
    if (saved[k] === undefined) delete process.env[k];
    else process.env[k] = saved[k];
  }
  mock.restoreAll();
});

function mockAxiosPost(responses) {
  let call = 0;
  mock.method(axios, 'post', async () => {
    const r = responses[call++] ?? responses[responses.length - 1];
    if (r instanceof Error) throw r;
    return r;
  });
}

function makeOkResponse(content = 'ok') {
  return { data: { choices: [{ message: { content } }] } };
}

describe('chat - key rotation', async () => {
  const { chat } = await import('./llm.js');

  it('succeeds on first key', async () => {
    process.env.OPENAI_API_KEY = 'key1';
    mockAxiosPost([makeOkResponse('hello')]);
    const result = await chat('sys', 'user');
    assert.equal(result, 'hello');
  });

  it('shuffles keys and tries each exactly once', async () => {
    process.env.OPENAI_API_KEY = 'key1,key2,key3';
    const usedKeys = [];
    mock.method(axios, 'post', async (_url, _body, opts) => {
      usedKeys.push(opts.headers['Authorization'].replace('Bearer ', ''));
      throw new Error('fail');
    });
    const exitMock = mock.method(process, 'exit', () => { throw new Error('exit:1'); });
    await assert.rejects(() => chat('sys', 'user'), /exit:1/);
    assert.equal(exitMock.mock.calls[0].arguments[0], 1);
    assert.equal(usedKeys.length, 3);
    assert.deepEqual([...usedKeys].sort(), ['key1', 'key2', 'key3']);
  });

  it('calls process.exit(1) when all keys fail', async () => {
    process.env.OPENAI_API_KEY = 'k1,k2';
    mockAxiosPost([new Error('fail'), new Error('fail')]);
    const exitMock = mock.method(process, 'exit', () => { throw new Error('exit:1'); });
    await assert.rejects(() => chat('sys', 'user'), /exit:1/);
    assert.equal(exitMock.mock.calls[0].arguments[0], 1);
  });

  it('does not set Authorization header for ollama', async () => {
    process.env.OLLAMA_BASE_URL = 'http://localhost:11434/v1';
    process.env.OLLAMA_MODEL = 'llama3';
    let capturedHeaders;
    mock.method(axios, 'post', async (_url, _body, opts) => {
      capturedHeaders = opts.headers;
      return makeOkResponse('ollama response');
    });
    await chat('sys', 'user');
    assert.equal(capturedHeaders['Authorization'], undefined);
  });

  it('sets Authorization header for openai', async () => {
    process.env.OPENAI_API_KEY = 'sk-test';
    let capturedHeaders;
    mock.method(axios, 'post', async (_url, _body, opts) => {
      capturedHeaders = opts.headers;
      return makeOkResponse();
    });
    await chat('sys', 'user');
    assert.equal(capturedHeaders['Authorization'], 'Bearer sk-test');
  });

  it('does not set timeout', async () => {
    process.env.OPENAI_API_KEY = 'sk-test';
    let capturedOpts;
    mock.method(axios, 'post', async (_url, _body, opts) => {
      capturedOpts = opts;
      return makeOkResponse();
    });
    await chat('sys', 'user');
    assert.equal(capturedOpts.timeout, undefined);
  });

  it('does not pass httpsAgent to axios', async () => {
    process.env.OPENAI_API_KEY = 'sk-test';
    let capturedOpts;
    mock.method(axios, 'post', async (_url, _body, opts) => {
      capturedOpts = opts;
      return makeOkResponse();
    });
    await chat('sys', 'user');
    assert.equal(capturedOpts.httpsAgent, undefined);
  });

  it('sets anthropic-version header for claude', async () => {
    process.env.CLAUDE_API_KEY = 'claude-key';
    let capturedHeaders;
    mock.method(axios, 'post', async (_url, _body, opts) => {
      capturedHeaders = opts.headers;
      return makeOkResponse();
    });
    await chat('sys', 'user');
    assert.equal(capturedHeaders['anthropic-version'], '2023-06-01');
  });
});

describe('chatJSON', async () => {
  const { chatJSON } = await import('./llm.js');

  it('parses plain JSON response', async () => {
    process.env.OPENAI_API_KEY = 'sk-test';
    mockAxiosPost([makeOkResponse('[{"level":"critical"}]')]);
    const result = await chatJSON('sys', 'user');
    assert.deepEqual(result, [{ level: 'critical' }]);
  });

  it('strips markdown code block before parsing', async () => {
    process.env.OPENAI_API_KEY = 'sk-test';
    mockAxiosPost([makeOkResponse('```json\n[{"level":"info"}]\n```')]);
    const result = await chatJSON('sys', 'user');
    assert.deepEqual(result, [{ level: 'info' }]);
  });

  it('returns [] when JSON is invalid', async () => {
    process.env.OPENAI_API_KEY = 'sk-test';
    mockAxiosPost([makeOkResponse('not json')]);
    const result = await chatJSON('sys', 'user');
    assert.deepEqual(result, []);
  });
});