Skip to content

Commit fc75398

Browse files
vinzeevinzee
committed
Add comprehensive tests for AI model configuration
This commit introduces a new test file to thoroughly validate the `getAIModel` function. The tests cover various scenarios including provider routing, API key handling, base URL configuration, model name selection, custom headers for OpenAI, and the `AI_EXTRA_BODY` functionality. This ensures robust and predictable behavior for AI model instantiation based on environment variables.
1 parent 3e05a19 commit fc75398

1 file changed

Lines changed: 399 additions & 0 deletions

File tree

‎packages/api/src/controllers/__tests__/ai.test.ts‎

Lines changed: 399 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,399 @@
1+
import type { LanguageModel } from 'ai';
2+
3+
const mockAnthropicModel = {
4+
modelId: 'claude-sonnet-4-5-20250929',
5+
} as unknown as LanguageModel;
6+
7+
const mockOpenAIModel = {
8+
modelId: 'gpt-4o',
9+
} as unknown as LanguageModel;
10+
11+
const mockAnthropicFactory = jest.fn(() => mockAnthropicModel);
12+
const mockCreateAnthropic = jest.fn(() => mockAnthropicFactory);
13+
14+
const mockOpenAIChatFactory = jest.fn(() => mockOpenAIModel);
15+
const mockCreateOpenAI = jest.fn(() => ({
16+
chat: mockOpenAIChatFactory,
17+
}));
18+
19+
jest.mock('@ai-sdk/anthropic', () => ({
20+
createAnthropic: (...args: unknown[]) => mockCreateAnthropic(...args),
21+
}));
22+
23+
jest.mock('@ai-sdk/openai', () => ({
24+
createOpenAI: (...args: unknown[]) => mockCreateOpenAI(...args),
25+
}));
26+
27+
jest.mock('@/utils/logger', () => ({
28+
__esModule: true,
29+
default: {
30+
info: jest.fn(),
31+
warn: jest.fn(),
32+
error: jest.fn(),
33+
},
34+
}));
35+
36+
const mockConfig: Record<string, unknown> = { __esModule: true };
37+
38+
jest.mock('@/config', () => mockConfig);
39+
40+
function setConfig(overrides: Record<string, string | undefined>) {
41+
Object.keys(mockConfig).forEach(k => {
42+
if (k !== '__esModule') delete mockConfig[k];
43+
});
44+
Object.assign(mockConfig, overrides);
45+
}
46+
47+
import { getAIModel } from '@/controllers/ai';
48+
49+
beforeEach(() => {
50+
setConfig({});
51+
jest.clearAllMocks();
52+
});
53+
54+
describe('getAIModel', () => {
55+
describe('provider routing', () => {
56+
it('throws when no provider is configured', () => {
57+
expect(() => getAIModel()).toThrow(
58+
'No AI provider configured. Set AI_PROVIDER and AI_API_KEY environment variables.',
59+
);
60+
});
61+
62+
it('throws on unknown provider', () => {
63+
setConfig({ AI_PROVIDER: 'gemini' });
64+
expect(() => getAIModel()).toThrow(
65+
'Unknown AI provider: gemini. Currently supported: anthropic, openai',
66+
);
67+
});
68+
69+
it('routes to anthropic when AI_PROVIDER=anthropic', () => {
70+
setConfig({
71+
AI_PROVIDER: 'anthropic',
72+
AI_API_KEY: 'sk-test',
73+
});
74+
const model = getAIModel();
75+
expect(model).toBe(mockAnthropicModel);
76+
expect(mockCreateAnthropic).toHaveBeenCalledTimes(1);
77+
});
78+
79+
it('routes to openai when AI_PROVIDER=openai', () => {
80+
setConfig({
81+
AI_PROVIDER: 'openai',
82+
AI_API_KEY: 'sk-test',
83+
AI_MODEL_NAME: 'gpt-4o',
84+
});
85+
const model = getAIModel();
86+
expect(model).toBe(mockOpenAIModel);
87+
expect(mockCreateOpenAI).toHaveBeenCalledTimes(1);
88+
});
89+
});
90+
91+
describe('legacy anthropic support', () => {
92+
it('falls back to anthropic when ANTHROPIC_API_KEY is set without AI_PROVIDER', () => {
93+
setConfig({
94+
ANTHROPIC_API_KEY: 'sk-ant-legacy',
95+
});
96+
const model = getAIModel();
97+
expect(model).toBe(mockAnthropicModel);
98+
expect(mockCreateAnthropic).toHaveBeenCalledWith(
99+
expect.objectContaining({ apiKey: 'sk-ant-legacy' }),
100+
);
101+
});
102+
});
103+
});
104+
105+
describe('anthropic provider', () => {
106+
it('throws when no API key is set', () => {
107+
setConfig({ AI_PROVIDER: 'anthropic' });
108+
expect(() => getAIModel()).toThrow(
109+
'No API key defined for Anthropic. Set AI_API_KEY or ANTHROPIC_API_KEY.',
110+
);
111+
});
112+
113+
it('uses AI_API_KEY over ANTHROPIC_API_KEY', () => {
114+
setConfig({
115+
AI_PROVIDER: 'anthropic',
116+
AI_API_KEY: 'sk-new',
117+
ANTHROPIC_API_KEY: 'sk-old',
118+
});
119+
getAIModel();
120+
expect(mockCreateAnthropic).toHaveBeenCalledWith(
121+
expect.objectContaining({ apiKey: 'sk-new' }),
122+
);
123+
});
124+
125+
it('passes baseURL when AI_BASE_URL is set', () => {
126+
setConfig({
127+
AI_PROVIDER: 'anthropic',
128+
AI_API_KEY: 'sk-test',
129+
AI_BASE_URL: 'https://custom.endpoint.com',
130+
});
131+
getAIModel();
132+
expect(mockCreateAnthropic).toHaveBeenCalledWith(
133+
expect.objectContaining({
134+
apiKey: 'sk-test',
135+
baseURL: 'https://custom.endpoint.com',
136+
}),
137+
);
138+
});
139+
140+
it('uses default model when AI_MODEL_NAME is not set', () => {
141+
setConfig({
142+
AI_PROVIDER: 'anthropic',
143+
AI_API_KEY: 'sk-test',
144+
});
145+
getAIModel();
146+
expect(mockAnthropicFactory).toHaveBeenCalledWith(
147+
'claude-sonnet-4-5-20250929',
148+
);
149+
});
150+
151+
it('uses custom model name when AI_MODEL_NAME is set', () => {
152+
setConfig({
153+
AI_PROVIDER: 'anthropic',
154+
AI_API_KEY: 'sk-test',
155+
AI_MODEL_NAME: 'claude-3-haiku-20240307',
156+
});
157+
getAIModel();
158+
expect(mockAnthropicFactory).toHaveBeenCalledWith(
159+
'claude-3-haiku-20240307',
160+
);
161+
});
162+
});
163+
164+
describe('openai provider', () => {
165+
it('throws when no API key is set', () => {
166+
setConfig({ AI_PROVIDER: 'openai' });
167+
expect(() => getAIModel()).toThrow(
168+
'No API key defined for OpenAI provider. Set AI_API_KEY.',
169+
);
170+
});
171+
172+
it('throws when no model name is set', () => {
173+
setConfig({
174+
AI_PROVIDER: 'openai',
175+
AI_API_KEY: 'sk-test',
176+
});
177+
expect(() => getAIModel()).toThrow(
178+
'No model name configured for OpenAI provider. Set AI_MODEL_NAME',
179+
);
180+
});
181+
182+
it('creates provider with minimal config', () => {
183+
setConfig({
184+
AI_PROVIDER: 'openai',
185+
AI_API_KEY: 'sk-test',
186+
AI_MODEL_NAME: 'gpt-4o',
187+
});
188+
getAIModel();
189+
expect(mockCreateOpenAI).toHaveBeenCalledWith(
190+
expect.objectContaining({ apiKey: 'sk-test' }),
191+
);
192+
expect(mockOpenAIChatFactory).toHaveBeenCalledWith('gpt-4o');
193+
});
194+
195+
it('passes baseURL when AI_BASE_URL is set', () => {
196+
setConfig({
197+
AI_PROVIDER: 'openai',
198+
AI_API_KEY: 'sk-test',
199+
AI_MODEL_NAME: 'gpt-4o',
200+
AI_BASE_URL: 'https://proxy.example.com/v1',
201+
});
202+
getAIModel();
203+
expect(mockCreateOpenAI).toHaveBeenCalledWith(
204+
expect.objectContaining({
205+
apiKey: 'sk-test',
206+
baseURL: 'https://proxy.example.com/v1',
207+
}),
208+
);
209+
});
210+
211+
describe('custom headers', () => {
212+
it('adds X-Client-Id header when AI_CLIENT_ID is set', () => {
213+
setConfig({
214+
AI_PROVIDER: 'openai',
215+
AI_API_KEY: 'sk-test',
216+
AI_MODEL_NAME: 'gpt-4o',
217+
AI_CLIENT_ID: 'MyApp',
218+
});
219+
getAIModel();
220+
expect(mockCreateOpenAI).toHaveBeenCalledWith(
221+
expect.objectContaining({
222+
headers: expect.objectContaining({
223+
'X-Client-Id': 'MyApp',
224+
}),
225+
}),
226+
);
227+
});
228+
229+
it('adds X-Username header when AI_USERNAME is set', () => {
230+
setConfig({
231+
AI_PROVIDER: 'openai',
232+
AI_API_KEY: 'sk-test',
233+
AI_MODEL_NAME: 'gpt-4o',
234+
AI_USERNAME: 'testuser',
235+
});
236+
getAIModel();
237+
expect(mockCreateOpenAI).toHaveBeenCalledWith(
238+
expect.objectContaining({
239+
headers: expect.objectContaining({
240+
'X-Username': 'testuser',
241+
}),
242+
}),
243+
);
244+
});
245+
246+
it('adds both headers when both are set', () => {
247+
setConfig({
248+
AI_PROVIDER: 'openai',
249+
AI_API_KEY: 'sk-test',
250+
AI_MODEL_NAME: 'gpt-4o',
251+
AI_CLIENT_ID: 'MyApp',
252+
AI_USERNAME: 'testuser',
253+
});
254+
getAIModel();
255+
expect(mockCreateOpenAI).toHaveBeenCalledWith(
256+
expect.objectContaining({
257+
headers: {
258+
'X-Client-Id': 'MyApp',
259+
'X-Username': 'testuser',
260+
},
261+
}),
262+
);
263+
});
264+
265+
it('omits headers object when neither is set', () => {
266+
setConfig({
267+
AI_PROVIDER: 'openai',
268+
AI_API_KEY: 'sk-test',
269+
AI_MODEL_NAME: 'gpt-4o',
270+
});
271+
getAIModel();
272+
const call = mockCreateOpenAI.mock.calls[0][0];
273+
expect(call.headers).toBeUndefined();
274+
});
275+
});
276+
277+
describe('AI_EXTRA_BODY', () => {
278+
it('throws when AI_EXTRA_BODY is not valid JSON', () => {
279+
setConfig({
280+
AI_PROVIDER: 'openai',
281+
AI_API_KEY: 'sk-test',
282+
AI_MODEL_NAME: 'gpt-4o',
283+
AI_EXTRA_BODY: '{bad json',
284+
});
285+
expect(() => getAIModel()).toThrow('AI_EXTRA_BODY is not valid JSON');
286+
});
287+
288+
it('provides a custom fetch when AI_EXTRA_BODY is valid JSON', () => {
289+
setConfig({
290+
AI_PROVIDER: 'openai',
291+
AI_API_KEY: 'sk-test',
292+
AI_MODEL_NAME: 'gpt-4o',
293+
AI_EXTRA_BODY: '{"extra_key":"extra_val"}',
294+
});
295+
getAIModel();
296+
const call = mockCreateOpenAI.mock.calls[0][0];
297+
expect(call.fetch).toBeDefined();
298+
expect(typeof call.fetch).toBe('function');
299+
});
300+
301+
it('does not provide custom fetch when AI_EXTRA_BODY is not set', () => {
302+
setConfig({
303+
AI_PROVIDER: 'openai',
304+
AI_API_KEY: 'sk-test',
305+
AI_MODEL_NAME: 'gpt-4o',
306+
});
307+
getAIModel();
308+
const call = mockCreateOpenAI.mock.calls[0][0];
309+
expect(call.fetch).toBeUndefined();
310+
});
311+
312+
it('custom fetch injects extra fields into JSON request body', async () => {
313+
const mockFetch = jest
314+
.spyOn(globalThis, 'fetch')
315+
.mockResolvedValue(new Response('{}'));
316+
317+
setConfig({
318+
AI_PROVIDER: 'openai',
319+
AI_API_KEY: 'sk-test',
320+
AI_MODEL_NAME: 'gpt-4o',
321+
AI_EXTRA_BODY: '{"extra_key":"extra_val","debug":true}',
322+
});
323+
getAIModel();
324+
const customFetch = mockCreateOpenAI.mock.calls[0][0].fetch;
325+
326+
await customFetch('https://api.example.com/v1/chat', {
327+
body: JSON.stringify({ model: 'gpt-4o', messages: [] }),
328+
method: 'POST',
329+
});
330+
331+
expect(mockFetch).toHaveBeenCalledTimes(1);
332+
const [, init] = mockFetch.mock.calls[0];
333+
const sentBody = JSON.parse(init!.body as string);
334+
expect(sentBody).toEqual({
335+
model: 'gpt-4o',
336+
messages: [],
337+
extra_key: 'extra_val',
338+
debug: true,
339+
});
340+
341+
mockFetch.mockRestore();
342+
});
343+
344+
it('custom fetch sends unmodified body when body is not JSON', async () => {
345+
const logger = jest.requireMock('@/utils/logger').default;
346+
const mockFetch = jest
347+
.spyOn(globalThis, 'fetch')
348+
.mockResolvedValue(new Response('{}'));
349+
350+
setConfig({
351+
AI_PROVIDER: 'openai',
352+
AI_API_KEY: 'sk-test',
353+
AI_MODEL_NAME: 'gpt-4o',
354+
AI_EXTRA_BODY: '{"extra_key":"extra_val"}',
355+
});
356+
getAIModel();
357+
const customFetch = mockCreateOpenAI.mock.calls[0][0].fetch;
358+
359+
await customFetch('https://api.example.com/v1/chat', {
360+
body: 'not-json',
361+
method: 'POST',
362+
});
363+
364+
expect(mockFetch).toHaveBeenCalledTimes(1);
365+
const [, init] = mockFetch.mock.calls[0];
366+
expect(init!.body).toBe('not-json');
367+
expect(logger.warn).toHaveBeenCalledWith(
368+
'AI_EXTRA_BODY: request body is not JSON, sending unmodified',
369+
);
370+
371+
mockFetch.mockRestore();
372+
});
373+
374+
it('custom fetch passes through when body is not a string', async () => {
375+
const mockFetch = jest
376+
.spyOn(globalThis, 'fetch')
377+
.mockResolvedValue(new Response('{}'));
378+
379+
setConfig({
380+
AI_PROVIDER: 'openai',
381+
AI_API_KEY: 'sk-test',
382+
AI_MODEL_NAME: 'gpt-4o',
383+
AI_EXTRA_BODY: '{"extra_key":"extra_val"}',
384+
});
385+
getAIModel();
386+
const customFetch = mockCreateOpenAI.mock.calls[0][0].fetch;
387+
388+
await customFetch('https://api.example.com/v1/chat', {
389+
method: 'GET',
390+
});
391+
392+
expect(mockFetch).toHaveBeenCalledTimes(1);
393+
const [, init] = mockFetch.mock.calls[0];
394+
expect(init!.body).toBeUndefined();
395+
396+
mockFetch.mockRestore();
397+
});
398+
});
399+
});

0 commit comments

Comments
 (0)