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); prisma = module.get(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" }, }); }); }); });