import { BadRequestException } from '@nestjs/common'; import { UserAiQuotaService } from './user-ai-quota.service'; describe('UserAiQuotaService', () => { let service: UserAiQuotaService; let mockFindUnique: jest.Mock; let mockUpsert: jest.Mock; const today = () => { const d = new Date(); return new Date(Date.UTC(d.getFullYear(), d.getMonth(), d.getDate())); }; beforeEach(() => { mockFindUnique = jest.fn(); mockUpsert = jest.fn(); const mockPrisma = { userAiSettings: { findUnique: mockFindUnique }, userAiUsageDaily: { findUnique: mockFindUnique, upsert: mockUpsert, }, } as any; service = new UserAiQuotaService(mockPrisma); }); // ═══════════════════════════════════════════════════════════════════ // checkQuota // ═══════════════════════════════════════════════════════════════════ describe('checkQuota', () => { it('passes when under both limits', async () => { mockFindUnique .mockResolvedValueOnce({ maxDailyAiJobs: 10, maxDailyTokenBudget: 50000 }) // settings .mockResolvedValueOnce({ jobCount: 3, totalTokens: 10000 }); // usage await expect(service.checkQuota('u1', 'platform_key')).resolves.toBeUndefined(); }); it('passes with defaults when settings is null', async () => { mockFindUnique .mockResolvedValueOnce(null) // no settings → defaults 20/100k .mockResolvedValueOnce({ jobCount: 5, totalTokens: 50000 }); await expect(service.checkQuota('u1', 'platform_key')).resolves.toBeUndefined(); }); it('passes when usage is null (first call of day)', async () => { mockFindUnique .mockResolvedValueOnce({ maxDailyAiJobs: 10, maxDailyTokenBudget: 50000 }) .mockResolvedValueOnce(null); // no usage yet → jobCount=0, totalTokens=0 await expect(service.checkQuota('u1', 'platform_key')).resolves.toBeUndefined(); }); it('throws DAILY_JOB_LIMIT_EXCEEDED when jobCount equals maxJobs', async () => { mockFindUnique .mockResolvedValueOnce({ maxDailyAiJobs: 10 }) .mockResolvedValueOnce({ jobCount: 10, totalTokens: 0 }); await expect(service.checkQuota('u1', 'platform_key')).rejects.toMatchObject({ response: expect.objectContaining({ errorCode: 'DAILY_JOB_LIMIT_EXCEEDED' }), }); }); it('throws DAILY_JOB_LIMIT_EXCEEDED when jobCount exceeds maxJobs', async () => { mockFindUnique .mockResolvedValueOnce({ maxDailyAiJobs: 5 }) .mockResolvedValueOnce({ jobCount: 8, totalTokens: 0 }); await expect(service.checkQuota('u1', 'platform_key')).rejects.toThrow(BadRequestException); }); it('throws DAILY_TOKEN_BUDGET_EXCEEDED when totalTokens >= maxTokens', async () => { mockFindUnique .mockResolvedValueOnce({ maxDailyTokenBudget: 10000 }) .mockResolvedValueOnce({ jobCount: 0, totalTokens: 10000 }); await expect(service.checkQuota('u1', 'user_deepseek_key')).rejects.toMatchObject({ response: expect.objectContaining({ errorCode: 'DAILY_TOKEN_BUDGET_EXCEEDED' }), }); }); it('uses separate quotas per apiKeyMode', async () => { // Check for platform_key mockFindUnique .mockResolvedValueOnce({ maxDailyAiJobs: 10 }) // settings .mockResolvedValueOnce(null); // platform_key usage: null → 0/0 await expect(service.checkQuota('u1', 'platform_key')).resolves.toBeUndefined(); // The usage query includes the correct apiKeyMode expect(mockFindUnique).toHaveBeenCalledWith( expect.objectContaining({ where: expect.objectContaining({ userId_localDate_apiKeyMode: expect.objectContaining({ apiKeyMode: 'platform_key', }), }), }), ); }); }); // ═══════════════════════════════════════════════════════════════════ // incrementJobCount // ═══════════════════════════════════════════════════════════════════ describe('incrementJobCount', () => { it('upserts with jobCount increment for today', async () => { const localDate = today(); mockUpsert.mockResolvedValue({}); await service.incrementJobCount('u1', 'platform_key'); expect(mockUpsert).toHaveBeenCalledWith({ where: { userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'platform_key' }, }, create: { userId: 'u1', localDate, apiKeyMode: 'platform_key', jobCount: 1 }, update: { jobCount: { increment: 1 } }, }); }); it('passes correct apiKeyMode for user key', async () => { mockUpsert.mockResolvedValue({}); await service.incrementJobCount('u2', 'user_deepseek_key'); const call = mockUpsert.mock.calls[0][0]; expect(call.where.userId_localDate_apiKeyMode.apiKeyMode).toBe('user_deepseek_key'); expect(call.create.jobCount).toBe(1); }); }); // ═══════════════════════════════════════════════════════════════════ // recordTokenUsage // ═══════════════════════════════════════════════════════════════════ describe('recordTokenUsage', () => { it('upserts with all token fields', async () => { const localDate = today(); mockUpsert.mockResolvedValue({}); await service.recordTokenUsage('u1', 'platform_key', 100, 50, 150, 0.002); expect(mockUpsert).toHaveBeenCalledWith({ where: { userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'platform_key' }, }, create: { userId: 'u1', localDate, apiKeyMode: 'platform_key', inputTokens: 100, outputTokens: 50, totalTokens: 150, costEstimate: 0.002, jobCount: 0, }, update: { inputTokens: { increment: 100 }, outputTokens: { increment: 50 }, totalTokens: { increment: 150 }, costEstimate: { increment: 0.002 }, }, }); }); it('defaults costEstimate to 0 when undefined', async () => { mockUpsert.mockResolvedValue({}); await service.recordTokenUsage('u1', 'platform_key', 100, 50, 150); const call = mockUpsert.mock.calls[0][0]; expect(call.create.costEstimate).toBe(0); expect(call.update.costEstimate).toEqual({ increment: 0 }); }); }); // ═══════════════════════════════════════════════════════════════════ // getUsage // ═══════════════════════════════════════════════════════════════════ describe('getUsage', () => { it('returns both platform and user key usage', async () => { mockFindUnique .mockResolvedValueOnce({ jobCount: 5, totalTokens: 30000 }) // platform_key .mockResolvedValueOnce({ jobCount: 2, totalTokens: 10000 }); // user_deepseek_key const result = await service.getUsage('u1'); expect(result).toEqual({ platformKey: { jobCount: 5, totalTokens: 30000 }, userDeepseekKey: { jobCount: 2, totalTokens: 10000 }, }); }); it('returns default empty objects when no records', async () => { mockFindUnique.mockResolvedValue(null); // both return null const result = await service.getUsage('u1'); expect(result).toEqual({ platformKey: { jobCount: 0, totalTokens: 0 }, userDeepseekKey: { jobCount: 0, totalTokens: 0 }, }); }); it('returns default for one mode when only platform has records', async () => { mockFindUnique .mockResolvedValueOnce({ jobCount: 3, totalTokens: 5000 }) .mockResolvedValueOnce(null); const result = await service.getUsage('u1'); expect(result.platformKey).toEqual({ jobCount: 3, totalTokens: 5000 }); expect(result.userDeepseekKey).toEqual({ jobCount: 0, totalTokens: 0 }); }); it('queries both modes in parallel', async () => { mockFindUnique.mockResolvedValue({ jobCount: 1, totalTokens: 100 }); await service.getUsage('u1'); const localDate = today(); expect(mockFindUnique).toHaveBeenCalledWith({ where: { userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'platform_key' } }, }); expect(mockFindUnique).toHaveBeenCalledWith({ where: { userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'user_deepseek_key' } }, }); }); }); });