api-server/src/modules/ai-runtime/user-ai-quota.service.spec.ts

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' } },
      });
    });
  });
});