test: add unit tests for platform-budget.service (API-AI-065)
All checks were successful
Deploy API Server / build-and-deploy (push) Successful in 45s
All checks were successful
Deploy API Server / build-and-deploy (push) Successful in 45s
- 15 tests covering all 6 public methods - checkPlatformBudget: healthy, auto-create, circuit-open, half-open limit, half-open pass, token exceeded, cost exceeded - recordSuccess: upsert with reset circuit to closed - recordFailure: increment failedCount, open circuit at threshold, first-day create - Admin: transitionToHalfOpen / closeCircuit - getBudgetState: state + limits, auto-create fallback Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
parent
c433b3dc5d
commit
4713758344
247
src/modules/ai-runtime/platform-budget.service.spec.ts
Normal file
247
src/modules/ai-runtime/platform-budget.service.spec.ts
Normal file
@ -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');
|
||||||
|
});
|
||||||
|
});
|
||||||
|
});
|
||||||
Loading…
x
Reference in New Issue
Block a user