From 9c7247aa9009f14066c585111ce5f277581e2561 Mon Sep 17 00:00:00 2001 From: wangdl Date: Thu, 18 Jun 2026 12:12:57 +0800 Subject: [PATCH] test: add unit tests for user-ai-quota.service (API-AI-064) - 15 tests covering all 4 public methods - checkQuota: pass under limits, defaults, null usage, job/token exceeded - incrementJobCount: correct upsert params, apiKeyMode routing - recordTokenUsage: all token fields, costEstimate default=0 - getUsage: both modes parallel, defaults for null records, mixed Co-Authored-By: Claude Opus 4.7 --- .../ai-runtime/user-ai-quota.service.spec.ts | 226 ++++++++++++++++++ 1 file changed, 226 insertions(+) create mode 100644 src/modules/ai-runtime/user-ai-quota.service.spec.ts diff --git a/src/modules/ai-runtime/user-ai-quota.service.spec.ts b/src/modules/ai-runtime/user-ai-quota.service.spec.ts new file mode 100644 index 0000000..99811f6 --- /dev/null +++ b/src/modules/ai-runtime/user-ai-quota.service.spec.ts @@ -0,0 +1,226 @@ +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' } }, + }); + }); + }); +});