wangdl/api-server
2
0
Fork 0
Code Issues 36 Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: split shared mockFindUnique into settings/usage mocks in quota tests
All checks were successful
Deploy API Server / build-and-deploy (push) Successful in 43s
Details

Browse Source 
- userAiSettings.findUnique → mockSettingsFindUnique
- userAiUsageDaily.findUnique → mockUsageFindUnique
- Eliminates order-dependency of mockResolvedValueOnce chains

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
wangdl 2026-06-18 12:14:33 +08:00
parent 9c7247aa90
commit c433b3dc5d
1 changed files with 36 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

87
src/modules/ai-runtime/user-ai-quota.service.spec.ts
View File

@ -3,7 +3,8 @@ import { UserAiQuotaService } from './user-ai-quota.service';
describe('UserAiQuotaService', () => { describe('UserAiQuotaService', () => {
let service: UserAiQuotaService; let service: UserAiQuotaService;
let mockFindUnique: jest.Mock; let mockSettingsFindUnique: jest.Mock;
let mockUsageFindUnique: jest.Mock;
let mockUpsert: jest.Mock; let mockUpsert: jest.Mock;
const today = () => { const today = () => {
@ -12,12 +13,13 @@ describe('UserAiQuotaService', () => {
}; };
beforeEach(() => { beforeEach(() => {
mockFindUnique = jest.fn(); mockSettingsFindUnique = jest.fn();
mockUsageFindUnique = jest.fn();
mockUpsert = jest.fn(); mockUpsert = jest.fn();
const mockPrisma = { const mockPrisma = {
userAiSettings: { findUnique: mockFindUnique }, userAiSettings: { findUnique: mockSettingsFindUnique },
userAiUsageDaily: { userAiUsageDaily: {
findUnique: mockFindUnique, findUnique: mockUsageFindUnique,
upsert: mockUpsert, upsert: mockUpsert,
}, },
} as any; } as any;
@ -30,33 +32,29 @@ describe('UserAiQuotaService', () => {
describe('checkQuota', () => { describe('checkQuota', () => {
it('passes when under both limits', async () => { it('passes when under both limits', async () => {
mockFindUnique mockSettingsFindUnique.mockResolvedValue({ maxDailyAiJobs: 10, maxDailyTokenBudget: 50000 });
.mockResolvedValueOnce({ maxDailyAiJobs: 10, maxDailyTokenBudget: 50000 }) // settings mockUsageFindUnique.mockResolvedValue({ jobCount: 3, totalTokens: 10000 });
.mockResolvedValueOnce({ jobCount: 3, totalTokens: 10000 }); // usage
await expect(service.checkQuota('u1', 'platform_key')).resolves.toBeUndefined(); await expect(service.checkQuota('u1', 'platform_key')).resolves.toBeUndefined();
}); });
it('passes with defaults when settings is null', async () => { it('passes with defaults when settings is null', async () => {
mockFindUnique mockSettingsFindUnique.mockResolvedValue(null);
.mockResolvedValueOnce(null) // no settings → defaults 20/100k mockUsageFindUnique.mockResolvedValue({ jobCount: 5, totalTokens: 50000 });
.mockResolvedValueOnce({ jobCount: 5, totalTokens: 50000 });
await expect(service.checkQuota('u1', 'platform_key')).resolves.toBeUndefined(); await expect(service.checkQuota('u1', 'platform_key')).resolves.toBeUndefined();
}); });
it('passes when usage is null (first call of day)', async () => { it('passes when usage is null (first call of day)', async () => {
mockFindUnique mockSettingsFindUnique.mockResolvedValue({ maxDailyAiJobs: 10, maxDailyTokenBudget: 50000 });
.mockResolvedValueOnce({ maxDailyAiJobs: 10, maxDailyTokenBudget: 50000 }) mockUsageFindUnique.mockResolvedValue(null); // no usage → defaults to 0
.mockResolvedValueOnce(null); // no usage yet → jobCount=0, totalTokens=0
await expect(service.checkQuota('u1', 'platform_key')).resolves.toBeUndefined(); await expect(service.checkQuota('u1', 'platform_key')).resolves.toBeUndefined();
}); });
it('throws DAILY_JOB_LIMIT_EXCEEDED when jobCount equals maxJobs', async () => { it('throws DAILY_JOB_LIMIT_EXCEEDED when jobCount equals maxJobs', async () => {
mockFindUnique mockSettingsFindUnique.mockResolvedValue({ maxDailyAiJobs: 10 });
.mockResolvedValueOnce({ maxDailyAiJobs: 10 }) mockUsageFindUnique.mockResolvedValue({ jobCount: 10, totalTokens: 0 });
.mockResolvedValueOnce({ jobCount: 10, totalTokens: 0 });
await expect(service.checkQuota('u1', 'platform_key')).rejects.toMatchObject({ await expect(service.checkQuota('u1', 'platform_key')).rejects.toMatchObject({
response: expect.objectContaining({ errorCode: 'DAILY_JOB_LIMIT_EXCEEDED' }), 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 () => { it('throws DAILY_JOB_LIMIT_EXCEEDED when jobCount exceeds maxJobs', async () => {
mockFindUnique mockSettingsFindUnique.mockResolvedValue({ maxDailyAiJobs: 5 });
.mockResolvedValueOnce({ maxDailyAiJobs: 5 }) mockUsageFindUnique.mockResolvedValue({ jobCount: 8, totalTokens: 0 });
.mockResolvedValueOnce({ jobCount: 8, totalTokens: 0 });
await expect(service.checkQuota('u1', 'platform_key')).rejects.toThrow(BadRequestException); await expect(service.checkQuota('u1', 'platform_key')).rejects.toThrow(BadRequestException);
}); });
it('throws DAILY_TOKEN_BUDGET_EXCEEDED when totalTokens >= maxTokens', async () => { it('throws DAILY_TOKEN_BUDGET_EXCEEDED when totalTokens >= maxTokens', async () => {
mockFindUnique mockSettingsFindUnique.mockResolvedValue({ maxDailyTokenBudget: 10000 });
.mockResolvedValueOnce({ maxDailyTokenBudget: 10000 }) mockUsageFindUnique.mockResolvedValue({ jobCount: 0, totalTokens: 10000 });
.mockResolvedValueOnce({ jobCount: 0, totalTokens: 10000 });
await expect(service.checkQuota('u1', 'user_deepseek_key')).rejects.toMatchObject({ await expect(service.checkQuota('u1', 'user_deepseek_key')).rejects.toMatchObject({
response: expect.objectContaining({ errorCode: 'DAILY_TOKEN_BUDGET_EXCEEDED' }), response: expect.objectContaining({ errorCode: 'DAILY_TOKEN_BUDGET_EXCEEDED' }),
}); });
}); });
it('uses separate quotas per apiKeyMode', async () => { it('queries usage with correct apiKeyMode', async () => {
// Check for platform_key mockSettingsFindUnique.mockResolvedValue({ maxDailyAiJobs: 10 });
mockFindUnique mockUsageFindUnique.mockResolvedValue(null);
.mockResolvedValueOnce({ maxDailyAiJobs: 10 }) // settings
.mockResolvedValueOnce(null); // platform_key usage: null → 0/0
await expect(service.checkQuota('u1', 'platform_key')).resolves.toBeUndefined(); await service.checkQuota('u1', 'platform_key');
// The usage query includes the correct apiKeyMode expect(mockUsageFindUnique).toHaveBeenCalledWith(
expect(mockFindUnique).toHaveBeenCalledWith(
expect.objectContaining({ expect.objectContaining({
where: expect.objectContaining({ where: expect.objectContaining({
userId_localDate_apiKeyMode: expect.objectContaining({ userId_localDate_apiKeyMode: expect.objectContaining({ apiKeyMode: 'platform_key' }),
apiKeyMode: 'platform_key',
}),
}), }),
}), }),
); );
@ -114,9 +105,7 @@ describe('UserAiQuotaService', () => {
await service.incrementJobCount('u1', 'platform_key'); await service.incrementJobCount('u1', 'platform_key');
expect(mockUpsert).toHaveBeenCalledWith({ expect(mockUpsert).toHaveBeenCalledWith({
where: { where: { userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'platform_key' } },
userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'platform_key' },
},
create: { userId: 'u1', localDate, apiKeyMode: 'platform_key', jobCount: 1 }, create: { userId: 'u1', localDate, apiKeyMode: 'platform_key', jobCount: 1 },
update: { jobCount: { increment: 1 } }, update: { jobCount: { increment: 1 } },
}); });
@ -125,9 +114,7 @@ describe('UserAiQuotaService', () => {
it('passes correct apiKeyMode for user key', async () => { it('passes correct apiKeyMode for user key', async () => {
mockUpsert.mockResolvedValue({}); mockUpsert.mockResolvedValue({});
await service.incrementJobCount('u2', 'user_deepseek_key'); await service.incrementJobCount('u2', 'user_deepseek_key');
const call = mockUpsert.mock.calls[0][0]; expect(mockUpsert.mock.calls[0][0].where.userId_localDate_apiKeyMode.apiKeyMode).toBe('user_deepseek_key');
expect(call.where.userId_localDate_apiKeyMode.apiKeyMode).toBe('user_deepseek_key');
expect(call.create.jobCount).toBe(1);
}); });
}); });
@ -143,9 +130,7 @@ describe('UserAiQuotaService', () => {
await service.recordTokenUsage('u1', 'platform_key', 100, 50, 150, 0.002); await service.recordTokenUsage('u1', 'platform_key', 100, 50, 150, 0.002);
expect(mockUpsert).toHaveBeenCalledWith({ expect(mockUpsert).toHaveBeenCalledWith({
where: { where: { userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'platform_key' } },
userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'platform_key' },
},
create: { create: {
userId: 'u1', localDate, apiKeyMode: 'platform_key', userId: 'u1', localDate, apiKeyMode: 'platform_key',
inputTokens: 100, outputTokens: 50, totalTokens: 150, inputTokens: 100, outputTokens: 50, totalTokens: 150,
@ -175,9 +160,9 @@ describe('UserAiQuotaService', () => {
describe('getUsage', () => { describe('getUsage', () => {
it('returns both platform and user key usage', async () => { it('returns both platform and user key usage', async () => {
mockFindUnique mockUsageFindUnique
.mockResolvedValueOnce({ jobCount: 5, totalTokens: 30000 }) // platform_key .mockResolvedValueOnce({ jobCount: 5, totalTokens: 30000 })
.mockResolvedValueOnce({ jobCount: 2, totalTokens: 10000 }); // user_deepseek_key .mockResolvedValueOnce({ jobCount: 2, totalTokens: 10000 });
const result = await service.getUsage('u1'); const result = await service.getUsage('u1');
@ -188,7 +173,7 @@ describe('UserAiQuotaService', () => {
}); });
it('returns default empty objects when no records', async () => { it('returns default empty objects when no records', async () => {
mockFindUnique.mockResolvedValue(null); // both return null mockUsageFindUnique.mockResolvedValue(null);
const result = await service.getUsage('u1'); const result = await service.getUsage('u1');
@ -198,8 +183,8 @@ describe('UserAiQuotaService', () => {
}); });
}); });
it('returns default for one mode when only platform has records', async () => { it('returns default for one mode when missing', async () => {
mockFindUnique mockUsageFindUnique
.mockResolvedValueOnce({ jobCount: 3, totalTokens: 5000 }) .mockResolvedValueOnce({ jobCount: 3, totalTokens: 5000 })
.mockResolvedValueOnce(null); .mockResolvedValueOnce(null);
@ -210,15 +195,15 @@ describe('UserAiQuotaService', () => {
}); });
it('queries both modes in parallel', async () => { it('queries both modes in parallel', async () => {
mockFindUnique.mockResolvedValue({ jobCount: 1, totalTokens: 100 }); mockUsageFindUnique.mockResolvedValue({ jobCount: 1, totalTokens: 100 });
await service.getUsage('u1'); await service.getUsage('u1');
const localDate = today(); const localDate = today();
expect(mockFindUnique).toHaveBeenCalledWith({ expect(mockUsageFindUnique).toHaveBeenCalledWith({
where: { userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'platform_key' } }, 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' } }, where: { userId_localDate_apiKeyMode: { userId: 'u1', localDate, apiKeyMode: 'user_deepseek_key' } },
}); });
}); });