129 lines
4.1 KiB
JavaScript
129 lines
4.1 KiB
JavaScript
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('rotates to second key when first fails', async () => {
|
|
process.env.OPENAI_API_KEY = 'key1,key2';
|
|
mockAxiosPost([new Error('rate limit'), makeOkResponse('from key2')]);
|
|
const result = await chat('sys', 'user');
|
|
assert.equal(result, 'from key2');
|
|
});
|
|
|
|
it('rotates through all keys and succeeds on last', async () => {
|
|
process.env.OPENAI_API_KEY = 'k1,k2,k3';
|
|
mockAxiosPost([new Error('fail'), new Error('fail'), makeOkResponse('from k3')]);
|
|
const result = await chat('sys', 'user');
|
|
assert.equal(result, 'from k3');
|
|
});
|
|
|
|
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('uses 30s 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, 30000);
|
|
});
|
|
|
|
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');
|
|
});
|
|
});
|