diff --git a/src/modules/ai-runtime/platform-budget.service.spec.ts b/src/modules/ai-runtime/platform-budget.service.spec.ts new file mode 100644 index 0000000..3b04b18 --- /dev/null +++ b/src/modules/ai-runtime/platform-budget.service.spec.ts @@ -0,0 +1,247 @@ +import { BadRequestException } from '@nestjs/common'; +import { PlatformBudgetService } from './platform-budget.service'; + +describe('PlatformBudgetService', () => { + let service: PlatformBudgetService; + let mockBudgetFindUnique: jest.Mock; + let mockBudgetCreate: jest.Mock; + let mockBudgetUpsert: jest.Mock; + let mockJobCount: jest.Mock; + + const healthyBudget = { + localDate: new Date(), + provider: 'deepseek', + model: 'deepseek-chat', + totalTokens: 1000, + costEstimate: 100, + jobCount: 5, + failedCount: 0, + circuitBreakerStatus: 'closed', + circuitBreakerReason: null, + }; + + const today = () => { + const d = new Date(); + return new Date(Date.UTC(d.getFullYear(), d.getMonth(), d.getDate())); + }; + + beforeEach(() => { + mockBudgetFindUnique = jest.fn(); + mockBudgetCreate = jest.fn(); + mockBudgetUpsert = jest.fn(); + mockJobCount = jest.fn(); + + const mockPrisma = { + platformAiBudgetDaily: { + findUnique: mockBudgetFindUnique, + create: mockBudgetCreate, + upsert: mockBudgetUpsert, + }, + aiRuntimeJob: { count: mockJobCount }, + } as any; + + service = new PlatformBudgetService(mockPrisma); + }); + + // ═══════════════════════════════════════════════════════════════════ + // checkPlatformBudget + // ═══════════════════════════════════════════════════════════════════ + + describe('checkPlatformBudget', () => { + it('passes when budget is healthy', async () => { + mockBudgetFindUnique.mockResolvedValue(healthyBudget); + + await expect(service.checkPlatformBudget('deepseek', 'deepseek-chat')).resolves.toBeUndefined(); + }); + + it('auto-creates budget on first call of day', async () => { + mockBudgetFindUnique.mockResolvedValue(null); + mockBudgetCreate.mockResolvedValue(healthyBudget); + + await expect(service.checkPlatformBudget('deepseek', 'deepseek-chat')).resolves.toBeUndefined(); + expect(mockBudgetCreate).toHaveBeenCalledWith( + expect.objectContaining({ data: expect.objectContaining({ provider: 'deepseek', model: 'deepseek-chat' }) }), + ); + }); + + it('throws PLATFORM_CIRCUIT_OPEN when circuit is open', async () => { + mockBudgetFindUnique.mockResolvedValue({ + ...healthyBudget, + circuitBreakerStatus: 'open', + circuitBreakerReason: 'Consecutive failures reached 10', + }); + + await expect(service.checkPlatformBudget('deepseek', 'deepseek-chat')).rejects.toMatchObject({ + response: expect.objectContaining({ errorCode: 'PLATFORM_CIRCUIT_OPEN' }), + }); + }); + + it('throws PLATFORM_CIRCUIT_HALF_OPEN when half_open + active jobs >= limit', async () => { + mockBudgetFindUnique.mockResolvedValue({ + ...healthyBudget, + circuitBreakerStatus: 'half_open', + }); + mockJobCount.mockResolvedValue(2); // halfOpenMaxJobs = 2 + + await expect(service.checkPlatformBudget('deepseek', 'deepseek-chat')).rejects.toMatchObject({ + response: expect.objectContaining({ errorCode: 'PLATFORM_CIRCUIT_HALF_OPEN' }), + }); + }); + + it('passes half_open when active jobs under limit', async () => { + mockBudgetFindUnique.mockResolvedValue({ + ...healthyBudget, + circuitBreakerStatus: 'half_open', + }); + mockJobCount.mockResolvedValue(1); // 1 < 2 (halfOpenMaxJobs) + + await expect(service.checkPlatformBudget('deepseek', 'deepseek-chat')).resolves.toBeUndefined(); + }); + + it('throws PLATFORM_TOKEN_BUDGET_EXCEEDED', async () => { + mockBudgetFindUnique.mockResolvedValue({ + ...healthyBudget, + totalTokens: 10_000_000, // at limit + }); + + await expect(service.checkPlatformBudget('deepseek', 'deepseek-chat')).rejects.toMatchObject({ + response: expect.objectContaining({ errorCode: 'PLATFORM_TOKEN_BUDGET_EXCEEDED' }), + }); + }); + + it('throws PLATFORM_COST_BUDGET_EXCEEDED', async () => { + mockBudgetFindUnique.mockResolvedValue({ + ...healthyBudget, + costEstimate: 50_000, // at limit ($500) + }); + + await expect(service.checkPlatformBudget('deepseek', 'deepseek-chat')).rejects.toMatchObject({ + response: expect.objectContaining({ errorCode: 'PLATFORM_COST_BUDGET_EXCEEDED' }), + }); + }); + }); + + // ═══════════════════════════════════════════════════════════════════ + // recordSuccess + // ═══════════════════════════════════════════════════════════════════ + + describe('recordSuccess', () => { + it('upserts with usage data and resets circuit to closed', async () => { + const localDate = today(); + mockBudgetUpsert.mockResolvedValue({}); + + await service.recordSuccess('deepseek', 'deepseek-chat', 100, 50, 150, 0.002); + + expect(mockBudgetUpsert).toHaveBeenCalledWith({ + where: { localDate_provider_model: { localDate, provider: 'deepseek', model: 'deepseek-chat' } }, + create: expect.objectContaining({ + inputTokens: 100, outputTokens: 50, totalTokens: 150, costEstimate: 0.002, + }), + update: expect.objectContaining({ + circuitBreakerStatus: 'closed', + circuitBreakerReason: null, + failedCount: 0, + }), + }); + }); + }); + + // ═══════════════════════════════════════════════════════════════════ + // recordFailure + // ═══════════════════════════════════════════════════════════════════ + + describe('recordFailure', () => { + it('increments failedCount', async () => { + mockBudgetFindUnique.mockResolvedValue(healthyBudget); // failedCount=0 + mockBudgetUpsert.mockResolvedValue({}); + + await service.recordFailure('deepseek', 'deepseek-chat', 'MODEL_TIMEOUT'); + + const call = mockBudgetUpsert.mock.calls[0][0]; + expect(call.update.failedCount).toBe(1); + }); + + it('opens circuit breaker when failedCount reaches threshold', async () => { + mockBudgetFindUnique.mockResolvedValue({ + ...healthyBudget, + failedCount: 9, // threshold is 10 → this failure = 10 + }); + mockBudgetUpsert.mockResolvedValue({}); + + await service.recordFailure('deepseek', 'deepseek-chat', 'MODEL_TIMEOUT'); + + const call = mockBudgetUpsert.mock.calls[0][0]; + expect(call.update.circuitBreakerStatus).toBe('open'); + expect(call.update.circuitBreakerReason).toContain('MODEL_TIMEOUT'); + }); + + it('creates budget on first failure of day', async () => { + mockBudgetFindUnique.mockResolvedValue(null); + mockBudgetCreate.mockResolvedValue(healthyBudget); + mockBudgetUpsert.mockResolvedValue({}); + + await service.recordFailure('deepseek', 'deepseek-chat', 'ERROR'); + + expect(mockBudgetUpsert).toHaveBeenCalled(); + }); + }); + + // ═══════════════════════════════════════════════════════════════════ + // transitionToHalfOpen / closeCircuit + // ═══════════════════════════════════════════════════════════════════ + + describe('admin circuit control', () => { + it('transitionToHalfOpen upserts with half_open status', async () => { + mockBudgetUpsert.mockResolvedValue({}); + await service.transitionToHalfOpen('deepseek', 'deepseek-chat'); + const call = mockBudgetUpsert.mock.calls[0][0]; + expect(call.update.circuitBreakerStatus).toBe('half_open'); + }); + + it('closeCircuit upserts with closed status + reset failedCount', async () => { + mockBudgetUpsert.mockResolvedValue({}); + await service.closeCircuit('deepseek', 'deepseek-chat'); + const call = mockBudgetUpsert.mock.calls[0][0]; + expect(call.update.circuitBreakerStatus).toBe('closed'); + expect(call.update.failedCount).toBe(0); + }); + }); + + // ═══════════════════════════════════════════════════════════════════ + // getBudgetState + // ═══════════════════════════════════════════════════════════════════ + + describe('getBudgetState', () => { + it('returns budget state with limits', async () => { + mockBudgetFindUnique.mockResolvedValue({ + ...healthyBudget, + provider: 'deepseek', + model: 'deepseek-chat', + totalTokens: 500000, + costEstimate: 25000, + jobCount: 12, + failedCount: 2, + circuitBreakerStatus: 'closed', + circuitBreakerReason: null, + }); + + const result = await service.getBudgetState('deepseek', 'deepseek-chat'); + + expect(result.provider).toBe('deepseek'); + expect(result.totalTokens).toBe(500000); + expect(result.costEstimateCents).toBe(25000); + expect(result.circuitBreakerStatus).toBe('closed'); + expect(result.limits.maxDailyTokens).toBe(10_000_000); + expect(result.limits.maxDailyCostCents).toBe(50_000); + }); + + it('auto-creates budget when none exists', async () => { + mockBudgetFindUnique.mockResolvedValue(null); + mockBudgetCreate.mockResolvedValue(healthyBudget); + + const result = await service.getBudgetState('deepseek', 'deepseek-chat'); + + expect(result.provider).toBe('deepseek'); + }); + }); +});