stack/apps/api/src/llm-usage/llm-usage.service.spec.ts

import { describe, it, expect, beforeEach, vi } from "vitest";
import { Test, TestingModule } from "@nestjs/testing";
import { LlmUsageService } from "./llm-usage.service";
import { PrismaService } from "../prisma/prisma.service";
import { TrackUsageDto, UsageAnalyticsQueryDto } from "./dto";

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

  const mockPrismaService = {
    llmUsageLog: {
      create: vi.fn(),
      findMany: vi.fn(),
      groupBy: vi.fn(),
      aggregate: vi.fn(),
    },
  };

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

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

    vi.clearAllMocks();
  });

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

  describe("trackUsage", () => {
    it("should create a new usage log entry", async () => {
      const trackUsageDto: TrackUsageDto = {
        workspaceId: "workspace-123",
        userId: "user-456",
        provider: "ollama",
        model: "llama3.2",
        promptTokens: 100,
        completionTokens: 50,
        totalTokens: 150,
        costCents: 0.15,
        taskType: "chat",
        durationMs: 1500,
      };

      const expectedResult = {
        id: "usage-789",
        ...trackUsageDto,
        createdAt: new Date(),
      };

      mockPrismaService.llmUsageLog.create.mockResolvedValue(expectedResult);

      const result = await service.trackUsage(trackUsageDto);

      expect(result).toEqual(expectedResult);
      expect(mockPrismaService.llmUsageLog.create).toHaveBeenCalledWith({
        data: trackUsageDto,
      });
    });

    it("should handle optional fields correctly", async () => {
      const trackUsageDto: TrackUsageDto = {
        workspaceId: "workspace-123",
        userId: "user-456",
        provider: "ollama",
        model: "llama3.2",
        promptTokens: 100,
        completionTokens: 50,
        totalTokens: 150,
      };

      const expectedResult = {
        id: "usage-789",
        ...trackUsageDto,
        costCents: null,
        taskType: null,
        durationMs: null,
        providerInstanceId: null,
        conversationId: null,
        createdAt: new Date(),
      };

      mockPrismaService.llmUsageLog.create.mockResolvedValue(expectedResult);

      const result = await service.trackUsage(trackUsageDto);

      expect(result).toEqual(expectedResult);
    });

    it("should throw error when database operation fails", async () => {
      const trackUsageDto: TrackUsageDto = {
        workspaceId: "workspace-123",
        userId: "user-456",
        provider: "ollama",
        model: "llama3.2",
        promptTokens: 100,
        completionTokens: 50,
        totalTokens: 150,
      };

      mockPrismaService.llmUsageLog.create.mockRejectedValue(new Error("Database error"));

      await expect(service.trackUsage(trackUsageDto)).rejects.toThrow("Database error");
    });
  });

  describe("getUsageAnalytics", () => {
    it("should return aggregated usage analytics", async () => {
      const query: UsageAnalyticsQueryDto = {
        workspaceId: "workspace-123",
      };

      const mockUsageLogs = [
        {
          provider: "ollama",
          model: "llama3.2",
          taskType: "chat",
          promptTokens: 100,
          completionTokens: 50,
          totalTokens: 150,
          costCents: 0.15,
          durationMs: 1500,
        },
        {
          provider: "ollama",
          model: "llama3.2",
          taskType: "embed",
          promptTokens: 200,
          completionTokens: 0,
          totalTokens: 200,
          costCents: 0.1,
          durationMs: 500,
        },
      ];

      mockPrismaService.llmUsageLog.findMany.mockResolvedValue(mockUsageLogs);

      const result = await service.getUsageAnalytics(query);

      expect(result.totalCalls).toBe(2);
      expect(result.totalPromptTokens).toBe(300);
      expect(result.totalCompletionTokens).toBe(50);
      expect(result.totalTokens).toBe(350);
      expect(result.totalCostCents).toBe(0.25);
      expect(result.averageDurationMs).toBe(1000);
      expect(result.byProvider).toHaveLength(1);
      expect(result.byModel).toHaveLength(1);
      expect(result.byTaskType).toHaveLength(2);
    });

    it("should filter by date range", async () => {
      const query: UsageAnalyticsQueryDto = {
        workspaceId: "workspace-123",
        startDate: "2024-01-01T00:00:00Z",
        endDate: "2024-01-31T23:59:59Z",
      };

      mockPrismaService.llmUsageLog.findMany.mockResolvedValue([]);

      await service.getUsageAnalytics(query);

      expect(mockPrismaService.llmUsageLog.findMany).toHaveBeenCalledWith({
        where: {
          workspaceId: "workspace-123",
          createdAt: {
            gte: new Date("2024-01-01T00:00:00Z"),
            lte: new Date("2024-01-31T23:59:59Z"),
          },
        },
      });
    });

    it("should filter by provider", async () => {
      const query: UsageAnalyticsQueryDto = {
        workspaceId: "workspace-123",
        provider: "ollama",
      };

      mockPrismaService.llmUsageLog.findMany.mockResolvedValue([]);

      await service.getUsageAnalytics(query);

      expect(mockPrismaService.llmUsageLog.findMany).toHaveBeenCalledWith({
        where: {
          workspaceId: "workspace-123",
          provider: "ollama",
        },
      });
    });

    it("should filter by model", async () => {
      const query: UsageAnalyticsQueryDto = {
        workspaceId: "workspace-123",
        model: "llama3.2",
      };

      mockPrismaService.llmUsageLog.findMany.mockResolvedValue([]);

      await service.getUsageAnalytics(query);

      expect(mockPrismaService.llmUsageLog.findMany).toHaveBeenCalledWith({
        where: {
          workspaceId: "workspace-123",
          model: "llama3.2",
        },
      });
    });

    it("should filter by userId", async () => {
      const query: UsageAnalyticsQueryDto = {
        workspaceId: "workspace-123",
        userId: "user-456",
      };

      mockPrismaService.llmUsageLog.findMany.mockResolvedValue([]);

      await service.getUsageAnalytics(query);

      expect(mockPrismaService.llmUsageLog.findMany).toHaveBeenCalledWith({
        where: {
          workspaceId: "workspace-123",
          userId: "user-456",
        },
      });
    });

    it("should handle empty results", async () => {
      const query: UsageAnalyticsQueryDto = {
        workspaceId: "workspace-123",
      };

      mockPrismaService.llmUsageLog.findMany.mockResolvedValue([]);

      const result = await service.getUsageAnalytics(query);

      expect(result.totalCalls).toBe(0);
      expect(result.totalPromptTokens).toBe(0);
      expect(result.totalCompletionTokens).toBe(0);
      expect(result.totalTokens).toBe(0);
      expect(result.totalCostCents).toBe(0);
      expect(result.averageDurationMs).toBe(0);
      expect(result.byProvider).toEqual([]);
      expect(result.byModel).toEqual([]);
      expect(result.byTaskType).toEqual([]);
    });

    it("should handle null values in aggregation", async () => {
      const query: UsageAnalyticsQueryDto = {
        workspaceId: "workspace-123",
      };

      const mockUsageLogs = [
        {
          provider: "ollama",
          model: "llama3.2",
          taskType: null,
          promptTokens: 100,
          completionTokens: 50,
          totalTokens: 150,
          costCents: null,
          durationMs: null,
        },
      ];

      mockPrismaService.llmUsageLog.findMany.mockResolvedValue(mockUsageLogs);

      const result = await service.getUsageAnalytics(query);

      expect(result.totalCalls).toBe(1);
      expect(result.totalCostCents).toBe(0);
      expect(result.averageDurationMs).toBe(0);
    });
  });

  describe("getUsageByWorkspace", () => {
    it("should return usage logs for a specific workspace", async () => {
      const workspaceId = "workspace-123";
      const mockLogs = [
        {
          id: "log-1",
          workspaceId,
          userId: "user-1",
          provider: "ollama",
          model: "llama3.2",
          promptTokens: 100,
          completionTokens: 50,
          totalTokens: 150,
          createdAt: new Date(),
        },
      ];

      mockPrismaService.llmUsageLog.findMany.mockResolvedValue(mockLogs);

      const result = await service.getUsageByWorkspace(workspaceId);

      expect(result).toEqual(mockLogs);
      expect(mockPrismaService.llmUsageLog.findMany).toHaveBeenCalledWith({
        where: { workspaceId },
        orderBy: { createdAt: "desc" },
      });
    });
  });

  describe("getUsageByProvider", () => {
    it("should return usage logs for a specific provider", async () => {
      const workspaceId = "workspace-123";
      const provider = "ollama";
      const mockLogs = [
        {
          id: "log-1",
          workspaceId,
          userId: "user-1",
          provider,
          model: "llama3.2",
          promptTokens: 100,
          completionTokens: 50,
          totalTokens: 150,
          createdAt: new Date(),
        },
      ];

      mockPrismaService.llmUsageLog.findMany.mockResolvedValue(mockLogs);

      const result = await service.getUsageByProvider(workspaceId, provider);

      expect(result).toEqual(mockLogs);
      expect(mockPrismaService.llmUsageLog.findMany).toHaveBeenCalledWith({
        where: { workspaceId, provider },
        orderBy: { createdAt: "desc" },
      });
    });
  });

  describe("getUsageByModel", () => {
    it("should return usage logs for a specific model", async () => {
      const workspaceId = "workspace-123";
      const model = "llama3.2";
      const mockLogs = [
        {
          id: "log-1",
          workspaceId,
          userId: "user-1",
          provider: "ollama",
          model,
          promptTokens: 100,
          completionTokens: 50,
          totalTokens: 150,
          createdAt: new Date(),
        },
      ];

      mockPrismaService.llmUsageLog.findMany.mockResolvedValue(mockLogs);

      const result = await service.getUsageByModel(workspaceId, model);

      expect(result).toEqual(mockLogs);
      expect(mockPrismaService.llmUsageLog.findMany).toHaveBeenCalledWith({
        where: { workspaceId, model },
        orderBy: { createdAt: "desc" },
      });
    });
  });
});