app/src/__tests__/ai.test.ts

import { describe, it, expect, beforeEach, beforeAll, afterAll } from "vitest";
import Fastify from "fastify";
import cookie from "@fastify/cookie";
import rateLimit from "@fastify/rate-limit";
import jwt from "jsonwebtoken";
import prisma from "../config/database";
import { config as appConfig } from "../config/env";
import { securityHeaders } from "../middleware/security";
import aiRoutes from "../routes/admin/ai";
import {
  computeCostUsd,
  recordUsage,
  getMonthSpendUsd,
  getBudgetUsd,
  assertBudgetAvailable,
  isConfigured,
} from "../services/ai.service";

const KIND = "test_ai"; // marker so we only clean our own rows

beforeEach(async () => {
  await prisma.ai_usage.deleteMany({ where: { kind: KIND } });
});
afterAll(async () => {
  await prisma.ai_usage.deleteMany({ where: { kind: KIND } });
});

describe("ai.service cost + budget", () => {
  it("computes Sonnet 4.6 cost from tokens", () => {
    // 1,000,000 input @ $3 + 1,000,000 output @ $15 = $18
    expect(
      computeCostUsd("claude-sonnet-4-6", 1_000_000, 1_000_000),
    ).toBeCloseTo(18, 6);
    expect(computeCostUsd("claude-sonnet-4-6", 4000, 500)).toBeCloseTo(
      0.0195,
      6,
    );
  });

  it("records usage with the computed cost", async () => {
    await recordUsage({
      userId: null,
      kind: KIND,
      model: "claude-sonnet-4-6",
      inputTokens: 4000,
      outputTokens: 500,
    });
    const rows = await prisma.ai_usage.findMany({ where: { kind: KIND } });
    expect(rows.length).toBe(1);
    expect(Number(rows[0].cost_usd)).toBeCloseTo(0.0195, 6);
  });

  it("sums only the current month's spend", async () => {
    await recordUsage({
      userId: null,
      kind: KIND,
      model: "claude-sonnet-4-6",
      inputTokens: 1_000_000,
      outputTokens: 0,
    }); // $3 this month
    // An old row (last year) must NOT count.
    await prisma.ai_usage.create({
      data: {
        kind: KIND,
        model: "claude-sonnet-4-6",
        input_tokens: 1_000_000,
        output_tokens: 0,
        cost_usd: 3,
        created_at: new Date("2000-01-01T00:00:00Z"),
      },
    });
    const spend = await getMonthSpendUsd();
    expect(spend).toBeGreaterThanOrEqual(3);
    expect(spend).toBeLessThan(6); // the year-2000 $3 is excluded
  });

  it("assertBudgetAvailable blocks at/over budget, allows under", async () => {
    const budget = await getBudgetUsd();
    // Under budget → null (allowed)
    expect(await assertBudgetAvailable()).toBeNull();
    // Push spend over budget for this month, then it must block with 402.
    await prisma.ai_usage.create({
      data: {
        kind: KIND,
        model: "claude-sonnet-4-6",
        input_tokens: 0,
        output_tokens: 0,
        cost_usd: budget + 1,
        created_at: new Date(),
      },
    });
    const blocked = await assertBudgetAvailable();
    expect(blocked).not.toBeNull();
    expect(blocked?.status).toBe(402);
  });

  it("isConfigured reflects the API key presence", () => {
    expect(typeof isConfigured()).toBe("boolean");
  });
});

describe("HTTP /api/admin/ai", () => {
  let app: Awaited<ReturnType<typeof buildAiApp>>;
  let adminToken: string;
  let noPermToken: string;
  let noPermRoleId: number;
  let noPermUserId: number;
  let savedKey: string;

  async function buildAiApp() {
    const a = Fastify({ logger: false });
    await a.register(cookie);
    await a.register(rateLimit, { max: 1000, timeWindow: "1 minute" });
    a.addHook("onRequest", securityHeaders);
    await a.register(aiRoutes, { prefix: "/api/admin/ai" });
    return a;
  }
  function token(user: {
    id: number;
    username: string;
    roleName: string | null;
  }) {
    return jwt.sign(
      { sub: user.id, username: user.username, role: user.roleName },
      appConfig.jwt.secret,
      { expiresIn: "15m" },
    );
  }

  beforeAll(async () => {
    savedKey = appConfig.anthropic.apiKey;
    app = await buildAiApp();
    const admin = await prisma.users.findFirst({
      where: { roles: { name: "admin" } },
      include: { roles: true },
    });
    if (!admin) throw new Error("admin not found");
    adminToken = token({
      id: admin.id,
      username: admin.username,
      roleName: admin.roles?.name ?? null,
    });
    const stamp = Date.now();
    const role = await prisma.roles.create({
      data: { name: `noperm_ai_${stamp}`, display_name: "No Perm AI" },
    });
    noPermRoleId = role.id;
    const u = await prisma.users.create({
      data: {
        username: `noperm_ai_${stamp}`,
        first_name: "No",
        last_name: "Perm",
        email: `noperm_ai_${stamp}@test.local`,
        password_hash:
          "$2a$10$invalidinvalidinvalidinvalidinvalidinvalidinvalidinvali",
        role_id: role.id,
        is_active: true,
      },
    });
    noPermUserId = u.id;
    noPermToken = token({
      id: u.id,
      username: u.username,
      roleName: role.name,
    });
  });
  afterAll(async () => {
    if (app) await app.close();
    (appConfig.anthropic as { apiKey: string }).apiKey = savedKey;
    if (noPermUserId)
      await prisma.users
        .deleteMany({ where: { id: noPermUserId } })
        .catch(() => {});
    if (noPermRoleId)
      await prisma.roles
        .deleteMany({ where: { id: noPermRoleId } })
        .catch(() => {});
  });

  it("GET /usage requires ai.use", async () => {
    const res = await app.inject({
      method: "GET",
      url: "/api/admin/ai/usage",
      headers: { Authorization: `Bearer ${noPermToken}` },
    });
    expect(res.statusCode).toBe(403);
  });

  it("GET /usage returns spend + budget for an admin", async () => {
    const res = await app.inject({
      method: "GET",
      url: "/api/admin/ai/usage",
      headers: { Authorization: `Bearer ${adminToken}` },
    });
    expect(res.statusCode).toBe(200);
    const body = res.json();
    expect(typeof body.data.budget_usd).toBe("number");
    expect(typeof body.data.month_spend_usd).toBe("number");
  });

  it("POST /chat returns 503 when AI is not configured", async () => {
    (appConfig.anthropic as { apiKey: string }).apiKey = "";
    const res = await app.inject({
      method: "POST",
      url: "/api/admin/ai/chat",
      headers: {
        Authorization: `Bearer ${adminToken}`,
        "Content-Type": "application/json",
      },
      payload: { messages: [{ role: "user", content: "ahoj" }] },
    });
    expect(res.statusCode).toBe(503);
    (appConfig.anthropic as { apiKey: string }).apiKey = savedKey;
  });
});