From c433b3dc5dd87d47329c0541dc5b13748b407cbb Mon Sep 17 00:00:00 2001 From: wangdl Date: Thu, 18 Jun 2026 12:14:33 +0800 Subject: [PATCH] refactor: split shared mockFindUnique into settings/usage mocks in quota tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - userAiSettings.findUnique → mockSettingsFindUnique - userAiUsageDaily.findUnique → mockUsageFindUnique - Eliminates order-dependency of mockResolvedValueOnce chains Co-Authored-By: Claude Opus 4.7 --- .../ai-runtime/user-ai-quota.service.spec.ts | 87 ++++++++----------- 1 file changed, 36 insertions(+), 51 deletions(-) diff --git a/src/modules/ai-runtime/user-ai-quota.service.spec.ts b/src/modules/ai-runtime/user-ai-quota.service.spec.ts index 99811f6..e7b1abc 100644 --- a/src/modules/ai-runtime/user-ai-quota.service.spec.ts +++ b/src/modules/ai-runtime/user-ai-quota.service.spec.ts @@ -3,7 +3,8 @@ import { UserAiQuotaService } from './user-ai-quota.service'; describe('UserAiQuotaService', () => { let service: UserAiQuotaService; - let mockFindUnique: jest.Mock; + let mockSettingsFindUnique: jest.Mock; + let mockUsageFindUnique: jest.Mock; let mockUpsert: jest.Mock; const today = () => { @@ -12,12 +13,13 @@ describe('UserAiQuotaService', () => { }; beforeEach(() => { - mockFindUnique = jest.fn(); + mockSettingsFindUnique = jest.fn(); + mockUsageFindUnique = jest.fn(); mockUpsert = jest.fn(); const mockPrisma = { - userAiSettings: { findUnique: mockFindUnique }, + userAiSettings: { findUnique: mockSettingsFindUnique }, userAiUsageDaily: { - findUnique: mockFindUnique, + findUnique: mockUsageFindUnique, upsert: mockUpsert, }, } as any; @@ -30,33 +32,29 @@ describe('UserAiQuotaService', () => { describe('checkQuota', () => { it('passes when under both limits', async () => { - mockFindUnique - .mockResolvedValueOnce({ maxDailyAiJobs: 10, maxDailyTokenBudget: 50000 }) // settings - .mockResolvedValueOnce({ jobCount: 3, totalTokens: 10000 }); // usage + mockSettingsFindUnique.mockResolvedValue({ maxDailyAiJobs: 10, maxDailyTokenBudget: 50000 }); + mockUsageFindUnique.mockResolvedValue({ jobCount: 3, totalTokens: 10000 }); 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 }); + mockSettingsFindUnique.mockResolvedValue(null); + mockUsageFindUnique.mockResolvedValue({ 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 + mockSettingsFindUnique.mockResolvedValue({ maxDailyAiJobs: 10, maxDailyTokenBudget: 50000 }); + mockUsageFindUnique.mockResolvedValue(null); // no usage → defaults to 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 }); + mockSettingsFindUnique.mockResolvedValue({ maxDailyAiJobs: 10 }); + mockUsageFindUnique.mockResolvedValue({ jobCount: 10, totalTokens: 0 }); await expect(service.checkQuota('u1', 'platform_key')).rejects.toMatchObject({ response: expect.objectContaining({ errorCode: 'DAILY_JOB_LIMIT_EXCEEDED' }), @@ -64,38 +62,31 @@ describe('UserAiQuotaService', () => { }); it('throws DAILY_JOB_LIMIT_EXCEEDED when jobCount exceeds maxJobs', async () => { - mockFindUnique - .mockResolvedValueOnce({ maxDailyAiJobs: 5 }) - .mockResolvedValueOnce({ jobCount: 8, totalTokens: 0 }); + mockSettingsFindUnique.mockResolvedValue({ maxDailyAiJobs: 5 }); + mockUsageFindUnique.mockResolvedValue({ 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 }); + mockSettingsFindUnique.mockResolvedValue({ maxDailyTokenBudget: 10000 }); + mockUsageFindUnique.mockResolvedValue({ 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 + it('queries usage with correct apiKeyMode', async () => { + mockSettingsFindUnique.mockResolvedValue({ maxDailyAiJobs: 10 }); + mockUsageFindUnique.mockResolvedValue(null); - await expect(service.checkQuota('u1', 'platform_key')).resolves.toBeUndefined(); + await service.checkQuota('u1', 'platform_key'); - // The usage query includes the correct apiKeyMode - expect(mockFindUnique).toHaveBeenCalledWith( + expect(mockUsageFindUnique).toHaveBeenCalledWith( expect.objectContaining({ where: expect.objectContaining({ - userId_localDate_apiKeyMode: expect.objectContaining({ - apiKeyMode: 'platform_key', - }), + userId_localDate_apiKeyMode: expect.objectContaining({ apiKeyMode: 'platform_key' }), }), }), ); @@ -114,9 +105,7 @@ describe('UserAiQuotaService', () => { await service.incrementJobCount('u1', 'platform_key'); expect(mockUpsert).toHaveBeenCalledWith({ - where: { - userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'platform_key' }, - }, + where: { userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'platform_key' } }, create: { userId: 'u1', localDate, apiKeyMode: 'platform_key', jobCount: 1 }, update: { jobCount: { increment: 1 } }, }); @@ -125,9 +114,7 @@ describe('UserAiQuotaService', () => { 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); + expect(mockUpsert.mock.calls[0][0].where.userId_localDate_apiKeyMode.apiKeyMode).toBe('user_deepseek_key'); }); }); @@ -143,9 +130,7 @@ describe('UserAiQuotaService', () => { await service.recordTokenUsage('u1', 'platform_key', 100, 50, 150, 0.002); expect(mockUpsert).toHaveBeenCalledWith({ - where: { - userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'platform_key' }, - }, + where: { userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'platform_key' } }, create: { userId: 'u1', localDate, apiKeyMode: 'platform_key', inputTokens: 100, outputTokens: 50, totalTokens: 150, @@ -175,9 +160,9 @@ describe('UserAiQuotaService', () => { 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 + mockUsageFindUnique + .mockResolvedValueOnce({ jobCount: 5, totalTokens: 30000 }) + .mockResolvedValueOnce({ jobCount: 2, totalTokens: 10000 }); const result = await service.getUsage('u1'); @@ -188,7 +173,7 @@ describe('UserAiQuotaService', () => { }); it('returns default empty objects when no records', async () => { - mockFindUnique.mockResolvedValue(null); // both return null + mockUsageFindUnique.mockResolvedValue(null); const result = await service.getUsage('u1'); @@ -198,8 +183,8 @@ describe('UserAiQuotaService', () => { }); }); - it('returns default for one mode when only platform has records', async () => { - mockFindUnique + it('returns default for one mode when missing', async () => { + mockUsageFindUnique .mockResolvedValueOnce({ jobCount: 3, totalTokens: 5000 }) .mockResolvedValueOnce(null); @@ -210,15 +195,15 @@ describe('UserAiQuotaService', () => { }); it('queries both modes in parallel', async () => { - mockFindUnique.mockResolvedValue({ jobCount: 1, totalTokens: 100 }); + mockUsageFindUnique.mockResolvedValue({ jobCount: 1, totalTokens: 100 }); await service.getUsage('u1'); const localDate = today(); - expect(mockFindUnique).toHaveBeenCalledWith({ + expect(mockUsageFindUnique).toHaveBeenCalledWith({ where: { userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'platform_key' } }, }); - expect(mockFindUnique).toHaveBeenCalledWith({ + expect(mockUsageFindUnique).toHaveBeenCalledWith({ where: { userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'user_deepseek_key' } }, }); });