stack/apps/api/src/token-budget/token-budget.service.spec.ts

import { describe, it, expect, beforeEach, vi } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { TokenBudgetService } from "./token-budget.service";
import { PrismaService } from "../prisma/prisma.service";
import { NotFoundException } from "@nestjs/common";
import type { TaskComplexity } from "./interfaces";
import { COMPLEXITY_BUDGETS } from "./interfaces";

describe("TokenBudgetService", () => {
  let service: TokenBudgetService;
  let prisma: PrismaService;

  const mockPrismaService = {
    tokenBudget: {
      create: vi.fn(),
      findUnique: vi.fn(),
      update: vi.fn(),
    },
  };

  const mockWorkspaceId = "550e8400-e29b-41d4-a716-446655440001";
  const mockTaskId = "550e8400-e29b-41d4-a716-446655440002";
  const mockAgentId = "test-agent-001";

  const mockTokenBudget = {
    id: "550e8400-e29b-41d4-a716-446655440003",
    taskId: mockTaskId,
    workspaceId: mockWorkspaceId,
    agentId: mockAgentId,
    allocatedTokens: 150000,
    estimatedComplexity: "medium" as TaskComplexity,
    inputTokensUsed: 50000,
    outputTokensUsed: 30000,
    totalTokensUsed: 80000,
    estimatedCost: null,
    startedAt: new Date("2026-01-31T10:00:00Z"),
    lastUpdatedAt: new Date("2026-01-31T10:30:00Z"),
    completedAt: null,
    budgetUtilization: 0.533,
    suspiciousPattern: false,
    suspiciousReason: null,
  };

  beforeEach(async () => {
    const module: TestingModule = await Test.createTestingModule({
      providers: [
        TokenBudgetService,
        {
          provide: PrismaService,
          useValue: mockPrismaService,
        },
      ],
    }).compile();

    service = module.get<TokenBudgetService>(TokenBudgetService);
    prisma = module.get<PrismaService>(PrismaService);

    vi.clearAllMocks();
  });

  it("should be defined", () => {
    expect(service).toBeDefined();
  });

  describe("allocateBudget", () => {
    it("should allocate budget for a new task", async () => {
      const allocateDto = {
        taskId: mockTaskId,
        workspaceId: mockWorkspaceId,
        agentId: mockAgentId,
        complexity: "medium" as TaskComplexity,
        allocatedTokens: 150000,
      };

      mockPrismaService.tokenBudget.create.mockResolvedValue(mockTokenBudget);

      const result = await service.allocateBudget(allocateDto);

      expect(result).toEqual(mockTokenBudget);
      expect(mockPrismaService.tokenBudget.create).toHaveBeenCalledWith({
        data: {
          taskId: allocateDto.taskId,
          workspaceId: allocateDto.workspaceId,
          agentId: allocateDto.agentId,
          allocatedTokens: allocateDto.allocatedTokens,
          estimatedComplexity: allocateDto.complexity,
        },
      });
    });
  });

  describe("updateUsage", () => {
    it("should update token usage and recalculate utilization", async () => {
      mockPrismaService.tokenBudget.findUnique.mockResolvedValue(mockTokenBudget);

      const updatedBudget = {
        ...mockTokenBudget,
        inputTokensUsed: 60000,
        outputTokensUsed: 40000,
        totalTokensUsed: 100000,
        budgetUtilization: 0.667,
      };

      mockPrismaService.tokenBudget.update.mockResolvedValue(updatedBudget);

      const result = await service.updateUsage(mockTaskId, 10000, 10000);

      expect(result).toEqual(updatedBudget);
      expect(mockPrismaService.tokenBudget.findUnique).toHaveBeenCalledWith({
        where: { taskId: mockTaskId },
      });
      expect(mockPrismaService.tokenBudget.update).toHaveBeenCalledWith({
        where: { taskId: mockTaskId },
        data: {
          inputTokensUsed: { increment: 10000 },
          outputTokensUsed: { increment: 10000 },
          totalTokensUsed: { increment: 20000 },
          budgetUtilization: expect.closeTo(0.667, 2),
        },
      });
    });

    it("should throw NotFoundException if budget does not exist", async () => {
      mockPrismaService.tokenBudget.findUnique.mockResolvedValue(null);

      await expect(service.updateUsage(mockTaskId, 1000, 1000)).rejects.toThrow(NotFoundException);
    });
  });

  describe("analyzeBudget", () => {
    it("should analyze budget and detect suspicious pattern for high remaining budget", async () => {
      mockPrismaService.tokenBudget.findUnique.mockResolvedValue(mockTokenBudget);

      const result = await service.analyzeBudget(mockTaskId);

      expect(result.taskId).toBe(mockTaskId);
      expect(result.allocatedTokens).toBe(150000);
      expect(result.usedTokens).toBe(80000);
      expect(result.remainingTokens).toBe(70000);
      expect(result.utilizationPercentage).toBeCloseTo(53.3, 1);
      // 46.7% remaining is suspicious (>20% threshold)
      expect(result.suspiciousPattern).toBe(true);
      expect(result.recommendation).toBe("review");
    });

    it("should not detect suspicious pattern when utilization is high", async () => {
      // 85% utilization (15% remaining - below 20% threshold)
      const highUtilizationBudget = {
        ...mockTokenBudget,
        inputTokensUsed: 65000,
        outputTokensUsed: 62500,
        totalTokensUsed: 127500,
        budgetUtilization: 0.85,
      };

      mockPrismaService.tokenBudget.findUnique.mockResolvedValue(highUtilizationBudget);

      const result = await service.analyzeBudget(mockTaskId);

      expect(result.utilizationPercentage).toBeCloseTo(85.0, 1);
      expect(result.suspiciousPattern).toBe(false);
      expect(result.recommendation).toBe("accept");
    });

    it("should throw NotFoundException if budget does not exist", async () => {
      mockPrismaService.tokenBudget.findUnique.mockResolvedValue(null);

      await expect(service.analyzeBudget(mockTaskId)).rejects.toThrow(NotFoundException);
    });
  });

  describe("checkSuspiciousDoneClaim", () => {
    it("should detect suspicious pattern when >20% budget remaining", async () => {
      // 30% budget remaining
      const budgetWithRemaining = {
        ...mockTokenBudget,
        inputTokensUsed: 50000,
        outputTokensUsed: 55000,
        totalTokensUsed: 105000,
        budgetUtilization: 0.7,
      };

      mockPrismaService.tokenBudget.findUnique.mockResolvedValue(budgetWithRemaining);

      const result = await service.checkSuspiciousDoneClaim(mockTaskId);

      expect(result.suspicious).toBe(true);
      expect(result.reason).toContain("30.0%");
    });

    it("should not flag as suspicious when <20% budget remaining", async () => {
      // 10% budget remaining
      const budgetNearlyDone = {
        ...mockTokenBudget,
        inputTokensUsed: 70000,
        outputTokensUsed: 65000,
        totalTokensUsed: 135000,
        budgetUtilization: 0.9,
      };

      mockPrismaService.tokenBudget.findUnique.mockResolvedValue(budgetNearlyDone);

      const result = await service.checkSuspiciousDoneClaim(mockTaskId);

      expect(result.suspicious).toBe(false);
      expect(result.reason).toBeUndefined();
    });

    it("should detect very low utilization (<10%)", async () => {
      // 5% utilization
      const budgetVeryLowUsage = {
        ...mockTokenBudget,
        inputTokensUsed: 4000,
        outputTokensUsed: 3500,
        totalTokensUsed: 7500,
        budgetUtilization: 0.05,
      };

      mockPrismaService.tokenBudget.findUnique.mockResolvedValue(budgetVeryLowUsage);

      const result = await service.checkSuspiciousDoneClaim(mockTaskId);

      expect(result.suspicious).toBe(true);
      expect(result.reason).toContain("Very low budget utilization");
    });
  });

  describe("getBudgetUtilization", () => {
    it("should return budget utilization percentage", async () => {
      mockPrismaService.tokenBudget.findUnique.mockResolvedValue(mockTokenBudget);

      const result = await service.getBudgetUtilization(mockTaskId);

      expect(result).toBeCloseTo(53.3, 1);
    });

    it("should throw NotFoundException if budget does not exist", async () => {
      mockPrismaService.tokenBudget.findUnique.mockResolvedValue(null);

      await expect(service.getBudgetUtilization(mockTaskId)).rejects.toThrow(NotFoundException);
    });
  });

  describe("markCompleted", () => {
    it("should mark budget as completed", async () => {
      mockPrismaService.tokenBudget.findUnique.mockResolvedValue(mockTokenBudget);

      const completedBudget = {
        ...mockTokenBudget,
        completedAt: new Date("2026-01-31T11:00:00Z"),
      };

      mockPrismaService.tokenBudget.update.mockResolvedValue(completedBudget);

      await service.markCompleted(mockTaskId);

      expect(mockPrismaService.tokenBudget.update).toHaveBeenCalledWith({
        where: { taskId: mockTaskId },
        data: {
          completedAt: expect.any(Date),
        },
      });
    });

    it("should throw NotFoundException if budget does not exist", async () => {
      mockPrismaService.tokenBudget.findUnique.mockResolvedValue(null);

      await expect(service.markCompleted(mockTaskId)).rejects.toThrow(NotFoundException);
    });
  });

  describe("getDefaultBudgetForComplexity", () => {
    it("should return correct budget for low complexity", () => {
      const result = service.getDefaultBudgetForComplexity("low");
      expect(result).toBe(COMPLEXITY_BUDGETS.low);
    });

    it("should return correct budget for medium complexity", () => {
      const result = service.getDefaultBudgetForComplexity("medium");
      expect(result).toBe(COMPLEXITY_BUDGETS.medium);
    });

    it("should return correct budget for high complexity", () => {
      const result = service.getDefaultBudgetForComplexity("high");
      expect(result).toBe(COMPLEXITY_BUDGETS.high);
    });

    it("should return correct budget for critical complexity", () => {
      const result = service.getDefaultBudgetForComplexity("critical");
      expect(result).toBe(COMPLEXITY_BUDGETS.critical);
    });
  });
});