/** * @license * Copyright 2025 AionUi (aionui.com) * SPDX-License-Identifier: Apache-2.0 */ import { describe, it, expect, beforeEach, vi } from 'vitest'; import { RotatingApiClient } from '@/common/api/RotatingApiClient'; import { AuthType } from '@office-ai/aioncli-core'; // Do not globally mock ApiKeyManager - we'll test integration with real ApiKeyManager // Only mock it selectively in specific tests that need custom behavior // Concrete test implementation of abstract RotatingApiClient type TestClient = { apiKey: string }; class TestRotatingApiClient extends RotatingApiClient { constructor(api_keys: string, authType: AuthType, options = {}) { super(api_keys, authType, (key) => ({ apiKey: key }), options); } } describe('RotatingApiClient', () => { beforeEach(() => { vi.clearAllMocks(); }); describe('initialization', () => { it('creates client for single key', () => { const client = new TestRotatingApiClient('sk-single', AuthType.USE_OPENAI); expect(client.hasMultipleKeys()).toBe(false); }); it('creates ApiKeyManager for comma-separated keys', () => { const client = new TestRotatingApiClient('key1,key2', AuthType.USE_OPENAI); expect(client.hasMultipleKeys()).toBe(true); }); it('creates ApiKeyManager for newline-separated keys', () => { const client = new TestRotatingApiClient('key1\nkey2', AuthType.USE_ANTHROPIC); expect(client.hasMultipleKeys()).toBe(true); }); it('does not create ApiKeyManager for single key', () => { const client = new TestRotatingApiClient('single-key', AuthType.USE_OPENAI); expect(client.hasMultipleKeys()).toBe(false); }); it('calls createClientFn with first key', () => { const createClientSpy = vi.fn((key) => ({ apiKey: key })); class SpyClient extends RotatingApiClient { constructor(api_keys: string) { super(api_keys, AuthType.USE_OPENAI, createClientSpy); } } new SpyClient('test-key'); expect(createClientSpy).toHaveBeenCalledWith('test-key'); }); }); describe('error handling', () => { it('identifies 401 as retryable', () => { const client = new TestRotatingApiClient('key', AuthType.USE_OPENAI); const error = { status: 401 }; expect((client as any).isRetryableError(error)).toBe(true); }); it('identifies 429 as retryable', () => { const client = new TestRotatingApiClient('key', AuthType.USE_OPENAI); const error = { status: 429 }; expect((client as any).isRetryableError(error)).toBe(true); }); it('identifies 503 as retryable', () => { const client = new TestRotatingApiClient('key', AuthType.USE_OPENAI); const error = { status: 503 }; expect((client as any).isRetryableError(error)).toBe(true); }); it('identifies 5xx as retryable', () => { const client = new TestRotatingApiClient('key', AuthType.USE_OPENAI); expect((client as any).isRetryableError({ status: 500 })).toBe(true); expect((client as any).isRetryableError({ status: 502 })).toBe(true); expect((client as any).isRetryableError({ status: 599 })).toBe(true); }); it('identifies 400 as non-retryable', () => { const client = new TestRotatingApiClient('key', AuthType.USE_OPENAI); const error = { status: 400 }; expect((client as any).isRetryableError(error)).toBe(false); }); it('identifies 404 as non-retryable', () => { const client = new TestRotatingApiClient('key', AuthType.USE_OPENAI); const error = { status: 404 }; expect((client as any).isRetryableError(error)).toBe(false); }); it('handles error with code instead of status', () => { const client = new TestRotatingApiClient('key', AuthType.USE_OPENAI); const error = { code: 429 }; expect((client as any).isRetryableError(error)).toBe(true); }); it('handles non-object error', () => { const client = new TestRotatingApiClient('key', AuthType.USE_OPENAI); expect((client as any).isRetryableError('string error')).toBe(false); expect((client as any).isRetryableError(null)).toBe(false); }); }); describe('executeWithRetry', () => { it('succeeds on first attempt', async () => { const client = new TestRotatingApiClient('key', AuthType.USE_OPENAI); const operation = vi.fn().mockResolvedValue('success'); const result = await client.executeWithRetry(operation); expect(result).toBe('success'); expect(operation).toHaveBeenCalledTimes(1); }); it('throws if no client initialized', async () => { const client = new TestRotatingApiClient('', AuthType.USE_OPENAI); const operation = vi.fn(); await expect(client.executeWithRetry(operation)).rejects.toThrow(/Client not initialized/); }); it('retries on retryable error', async () => { const client = new TestRotatingApiClient('key', AuthType.USE_OPENAI, { maxRetries: 3 }); const operation = vi.fn().mockRejectedValueOnce({ status: 429 }).mockResolvedValue('success'); const result = await client.executeWithRetry(operation); expect(result).toBe('success'); expect(operation).toHaveBeenCalledTimes(2); }); it('exhausts retries and throws last error', async () => { const client = new TestRotatingApiClient('key', AuthType.USE_OPENAI, { maxRetries: 2 }); const operation = vi.fn().mockRejectedValue({ status: 503 }); await expect(client.executeWithRetry(operation)).rejects.toEqual({ status: 503 }); expect(operation).toHaveBeenCalledTimes(2); }); it('stops retrying on non-retryable error', async () => { const client = new TestRotatingApiClient('key', AuthType.USE_OPENAI, { maxRetries: 3 }); const operation = vi.fn().mockRejectedValue({ status: 400 }); await expect(client.executeWithRetry(operation)).rejects.toEqual({ status: 400 }); expect(operation).toHaveBeenCalledTimes(1); }); it('rotates key on retryable error with multiple keys (integration test)', async () => { // Use real ApiKeyManager with multiple keys const client = new TestRotatingApiClient('key1,key2,key3', AuthType.USE_OPENAI, { maxRetries: 3 }); let attemptCount = 0; const operation = vi.fn().mockImplementation(() => { attemptCount++; if (attemptCount === 1) { return Promise.reject({ status: 401 }); } return Promise.resolve('success'); }); const result = await client.executeWithRetry(operation); expect(result).toBe('success'); expect(operation).toHaveBeenCalledTimes(2); // Key should have rotated const status = client.getKeyStatus(); expect(status).toBeTruthy(); expect(status.blacklisted.length).toBeGreaterThan(0); }); it('does not rotate key on last attempt', async () => { // Use real ApiKeyManager with multiple keys const client = new TestRotatingApiClient('key1,key2', AuthType.USE_OPENAI, { maxRetries: 1 }); const operation = vi.fn().mockRejectedValue({ status: 429 }); await expect(client.executeWithRetry(operation)).rejects.toEqual({ status: 429 }); expect(operation).toHaveBeenCalledTimes(1); // No rotation should happen on last attempt const status = client.getKeyStatus(); expect(status).toBeTruthy(); // Blacklist should still be empty because it's the last attempt expect(status.blacklisted.length).toBe(0); }); }); describe('options', () => { it('uses default maxRetries and retryDelay', () => { const client = new TestRotatingApiClient('key', AuthType.USE_OPENAI); expect((client as any).options.maxRetries).toBe(3); expect((client as any).options.retryDelay).toBe(1000); }); it('accepts custom maxRetries', () => { const client = new TestRotatingApiClient('key', AuthType.USE_OPENAI, { maxRetries: 5 }); expect((client as any).options.maxRetries).toBe(5); }); it('accepts custom retryDelay', () => { const client = new TestRotatingApiClient('key', AuthType.USE_OPENAI, { retryDelay: 2000 }); expect((client as any).options.retryDelay).toBe(2000); }); }); });