ai-switcher/src/__tests__/AIClient.test.ts

A package which helps you switch between AI APIs using configurations, so that code changes are not required.

// src/__tests__/AIClient.test.ts
import { AIClient } from '../index';
import { AnthropicClient } from '../clients/anthropic';
import { OpenAIClient } from '../clients/openai';
import { Message, AIError } from '../types';
import { mockOpenAIClient, mockAnthropicClient } from './mocks';

jest.mock('../clients/anthropic', () => ({
  AnthropicClient: jest.fn().mockImplementation(() => mockAnthropicClient)
}));

jest.mock('../clients/openai', () => ({
  OpenAIClient: jest.fn().mockImplementation(() => mockOpenAIClient)
}));

describe('AIClient', () => {
  let client: AIClient;

  beforeEach(() => {
    jest.clearAllMocks();
    client = new AIClient({
      anthropicApiKey: 'test-anthropic-key',
      openaiApiKey: 'test-openai-key',
      defaultProvider: 'openai'
    });
  });

  describe('initialization', () => {
    it('should initialize with API keys', () => {
      expect(client).toBeInstanceOf(AIClient);
    });

    it('should throw error when no API keys provided', () => {
      expect(() => new AIClient({})).toThrow('At least one provider API key must be provided');
    });
  });

  describe('provider selection', () => {
    const messages: Message[] = [{ role: 'user' as const, content: 'test' }];

    it('should use default provider when none specified', async () => {
      mockOpenAIClient.createCompletion.mockResolvedValue('test response');

      await client.createCompletion(messages);
      expect(mockOpenAIClient.createCompletion).toHaveBeenCalled();
      expect(mockAnthropicClient.createCompletion).not.toHaveBeenCalled();
    });

    it('should override default provider when specified in options', async () => {
      mockAnthropicClient.createCompletion.mockResolvedValue('test response');

      await client.createCompletion(messages, { provider: 'anthropic' });
      expect(mockAnthropicClient.createCompletion).toHaveBeenCalled();
      expect(mockOpenAIClient.createCompletion).not.toHaveBeenCalled();
    });
  });

  describe('createEmbedding', () => {
    it('should return embedding data using openai provider', async () => {
      mockOpenAIClient.createEmbedding.mockResolvedValue([0.1, 0.2, 0.3]);
      const result = await client.createEmbedding('test', 'openai');
      expect(result).toEqual([0.1, 0.2, 0.3]);
    });

    it('should throw error for anthropic provider', async () => {
      mockAnthropicClient.createEmbedding.mockRejectedValue(new Error('Embeddings not supported by Anthropic'));
      await expect(client.createEmbedding('test', 'anthropic')).rejects.toThrow('Embeddings not supported by Anthropic');
    });
  });

  describe('error handling', () => {
    const messages: Message[] = [{ role: 'user' as const, content: 'test' }];

    it('should handle rate limit errors', async () => {
      const error = new AIError('Rate limit exceeded', 'openai', 429);
      mockOpenAIClient.createCompletion.mockRejectedValue(error);

      await expect(client.createCompletion(messages)).rejects.toMatchObject({
        message: expect.stringContaining('Rate limit exceeded'),
        provider: 'openai',
        statusCode: 429
      });
    });

    it('should handle authentication errors', async () => {
      const error = new AIError('Invalid API key', 'openai', 401);
      mockOpenAIClient.createCompletion.mockRejectedValue(error);

      await expect(client.createCompletion(messages)).rejects.toMatchObject({
        message: expect.stringContaining('Invalid API key'),
        provider: 'openai',
        statusCode: 401
      });
    });
  });
});