Add test coverage

This commit is contained in:
Matt Rubens
2025-01-07 01:52:42 -05:00
parent 537514de44
commit 5fbfe9b775
19 changed files with 3106 additions and 493 deletions

View File

@@ -1,167 +1,203 @@
import { DeepSeekHandler } from '../deepseek'
import { ApiHandlerOptions } from '../../../shared/api'
import OpenAI from 'openai'
import { Anthropic } from '@anthropic-ai/sdk'
import { DeepSeekHandler } from '../deepseek';
import { ApiHandlerOptions, deepSeekDefaultModelId } from '../../../shared/api';
import OpenAI from 'openai';
import { Anthropic } from '@anthropic-ai/sdk';
// Mock dependencies
jest.mock('openai')
// Mock OpenAI client
const mockCreate = jest.fn();
jest.mock('openai', () => {
return {
__esModule: true,
default: jest.fn().mockImplementation(() => ({
chat: {
completions: {
create: mockCreate.mockImplementation(async (options) => {
if (!options.stream) {
return {
id: 'test-completion',
choices: [{
message: { role: 'assistant', content: 'Test response', refusal: null },
finish_reason: 'stop',
index: 0
}],
usage: {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15
}
};
}
// Return async iterator for streaming
return {
[Symbol.asyncIterator]: async function* () {
yield {
choices: [{
delta: { content: 'Test response' },
index: 0
}],
usage: null
};
yield {
choices: [{
delta: {},
index: 0
}],
usage: {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15
}
};
}
};
})
}
}
}))
};
});
describe('DeepSeekHandler', () => {
const mockOptions: ApiHandlerOptions = {
deepSeekApiKey: 'test-key',
deepSeekModelId: 'deepseek-chat',
}
let handler: DeepSeekHandler;
let mockOptions: ApiHandlerOptions;
beforeEach(() => {
jest.clearAllMocks()
})
mockOptions = {
deepSeekApiKey: 'test-api-key',
deepSeekModelId: 'deepseek-chat',
deepSeekBaseUrl: 'https://api.deepseek.com/v1'
};
handler = new DeepSeekHandler(mockOptions);
mockCreate.mockClear();
});
test('constructor initializes with correct options', () => {
const handler = new DeepSeekHandler(mockOptions)
expect(handler).toBeInstanceOf(DeepSeekHandler)
expect(OpenAI).toHaveBeenCalledWith({
baseURL: 'https://api.deepseek.com/v1',
apiKey: mockOptions.deepSeekApiKey,
})
})
describe('constructor', () => {
it('should initialize with provided options', () => {
expect(handler).toBeInstanceOf(DeepSeekHandler);
expect(handler.getModel().id).toBe(mockOptions.deepSeekModelId);
});
test('getModel returns correct model info', () => {
const handler = new DeepSeekHandler(mockOptions)
const result = handler.getModel()
expect(result).toEqual({
id: mockOptions.deepSeekModelId,
info: expect.objectContaining({
maxTokens: 8192,
contextWindow: 64000,
supportsPromptCache: false,
supportsImages: false,
inputPrice: 0.014,
outputPrice: 0.28,
})
})
})
it('should throw error if API key is missing', () => {
expect(() => {
new DeepSeekHandler({
...mockOptions,
deepSeekApiKey: undefined
});
}).toThrow('DeepSeek API key is required');
});
test('getModel returns default model info when no model specified', () => {
const handler = new DeepSeekHandler({ deepSeekApiKey: 'test-key' })
const result = handler.getModel()
expect(result.id).toBe('deepseek-chat')
expect(result.info.maxTokens).toBe(8192)
})
it('should use default model ID if not provided', () => {
const handlerWithoutModel = new DeepSeekHandler({
...mockOptions,
deepSeekModelId: undefined
});
expect(handlerWithoutModel.getModel().id).toBe(deepSeekDefaultModelId);
});
test('createMessage handles string content correctly', async () => {
const handler = new DeepSeekHandler(mockOptions)
const mockStream = {
async *[Symbol.asyncIterator]() {
yield {
choices: [{
delta: {
content: 'test response'
}
}]
}
}
}
it('should use default base URL if not provided', () => {
const handlerWithoutBaseUrl = new DeepSeekHandler({
...mockOptions,
deepSeekBaseUrl: undefined
});
expect(handlerWithoutBaseUrl).toBeInstanceOf(DeepSeekHandler);
// The base URL is passed to OpenAI client internally
expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({
baseURL: 'https://api.deepseek.com/v1'
}));
});
const mockCreate = jest.fn().mockResolvedValue(mockStream)
;(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
completions: { create: mockCreate }
} as any
it('should use custom base URL if provided', () => {
const customBaseUrl = 'https://custom.deepseek.com/v1';
const handlerWithCustomUrl = new DeepSeekHandler({
...mockOptions,
deepSeekBaseUrl: customBaseUrl
});
expect(handlerWithCustomUrl).toBeInstanceOf(DeepSeekHandler);
// The custom base URL is passed to OpenAI client
expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({
baseURL: customBaseUrl
}));
});
const systemPrompt = 'test system prompt'
const messages: Anthropic.Messages.MessageParam[] = [
{ role: 'user', content: 'test message' }
]
it('should set includeMaxTokens to true', () => {
// Create a new handler and verify OpenAI client was called with includeMaxTokens
new DeepSeekHandler(mockOptions);
expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({
apiKey: mockOptions.deepSeekApiKey
}));
});
});
const generator = handler.createMessage(systemPrompt, messages)
const chunks = []
for await (const chunk of generator) {
chunks.push(chunk)
}
describe('getModel', () => {
it('should return model info for valid model ID', () => {
const model = handler.getModel();
expect(model.id).toBe(mockOptions.deepSeekModelId);
expect(model.info).toBeDefined();
expect(model.info.maxTokens).toBe(8192);
expect(model.info.contextWindow).toBe(64_000);
expect(model.info.supportsImages).toBe(false);
expect(model.info.supportsPromptCache).toBe(false);
});
expect(chunks).toHaveLength(1)
expect(chunks[0]).toEqual({
type: 'text',
text: 'test response'
})
it('should return provided model ID with default model info if model does not exist', () => {
const handlerWithInvalidModel = new DeepSeekHandler({
...mockOptions,
deepSeekModelId: 'invalid-model'
});
const model = handlerWithInvalidModel.getModel();
expect(model.id).toBe('invalid-model'); // Returns provided ID
expect(model.info).toBeDefined();
expect(model.info).toBe(handler.getModel().info); // But uses default model info
});
expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({
model: mockOptions.deepSeekModelId,
messages: [
{ role: 'system', content: systemPrompt },
{ role: 'user', content: 'test message' }
],
temperature: 0,
stream: true,
max_tokens: 8192,
stream_options: { include_usage: true }
}))
})
it('should return default model if no model ID is provided', () => {
const handlerWithoutModel = new DeepSeekHandler({
...mockOptions,
deepSeekModelId: undefined
});
const model = handlerWithoutModel.getModel();
expect(model.id).toBe(deepSeekDefaultModelId);
expect(model.info).toBeDefined();
});
});
test('createMessage handles complex content correctly', async () => {
const handler = new DeepSeekHandler(mockOptions)
const mockStream = {
async *[Symbol.asyncIterator]() {
yield {
choices: [{
delta: {
content: 'test response'
}
}]
}
}
}
const mockCreate = jest.fn().mockResolvedValue(mockStream)
;(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
completions: { create: mockCreate }
} as any
const systemPrompt = 'test system prompt'
describe('createMessage', () => {
const systemPrompt = 'You are a helpful assistant.';
const messages: Anthropic.Messages.MessageParam[] = [
{
role: 'user',
content: [
{ type: 'text', text: 'part 1' },
{ type: 'text', text: 'part 2' }
]
content: [{
type: 'text' as const,
text: 'Hello!'
}]
}
]
];
const generator = handler.createMessage(systemPrompt, messages)
await generator.next()
expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({
messages: [
{ role: 'system', content: systemPrompt },
{
role: 'user',
content: [
{ type: 'text', text: 'part 1' },
{ type: 'text', text: 'part 2' }
]
}
]
}))
})
test('createMessage handles API errors', async () => {
const handler = new DeepSeekHandler(mockOptions)
const mockStream = {
async *[Symbol.asyncIterator]() {
throw new Error('API Error')
it('should handle streaming responses', async () => {
const stream = handler.createMessage(systemPrompt, messages);
const chunks: any[] = [];
for await (const chunk of stream) {
chunks.push(chunk);
}
}
const mockCreate = jest.fn().mockResolvedValue(mockStream)
;(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
completions: { create: mockCreate }
} as any
expect(chunks.length).toBeGreaterThan(0);
const textChunks = chunks.filter(chunk => chunk.type === 'text');
expect(textChunks).toHaveLength(1);
expect(textChunks[0].text).toBe('Test response');
});
const generator = handler.createMessage('test', [])
await expect(generator.next()).rejects.toThrow('API Error')
})
})
it('should include usage information', async () => {
const stream = handler.createMessage(systemPrompt, messages);
const chunks: any[] = [];
for await (const chunk of stream) {
chunks.push(chunk);
}
const usageChunks = chunks.filter(chunk => chunk.type === 'usage');
expect(usageChunks.length).toBeGreaterThan(0);
expect(usageChunks[0].inputTokens).toBe(10);
expect(usageChunks[0].outputTokens).toBe(5);
});
});
});