app/src/__tests__/ai-budget-permission.test.ts

import { describe, it, expect, 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 } from "../config/env";
import { securityHeaders } from "../middleware/security";
import aiRoutes from "../routes/admin/ai";
import { getBudgetUsd, setBudgetUsd } from "../services/ai.service";

/**
 * Pinning test for audit finding L8a: the company-wide AI monthly budget must
 * NOT be mutable by an ordinary ai.use holder (budget_usd=0 is a company-wide
 * AI DoS; a huge value removes the cost cap). It is a company setting —
 * settings.company / settings.system / admin only.
 */

const N = "ai_budget_perm_";
const stamp = Date.now().toString(36);

let app: ReturnType<typeof Fastify>;
let adminToken: string;
let aiUserToken: string;
let aiRoleId: number;
let savedBudget: number;

function token(user: { id: number; username: string; roleName: string }) {
  return jwt.sign(
    { sub: user.id, username: user.username, role: user.roleName },
    config.jwt.secret,
    { expiresIn: "15m" },
  );
}

async function cleanup() {
  await prisma.users.deleteMany({ where: { username: { startsWith: N } } });
  const roles = await prisma.roles.findMany({
    where: { name: { startsWith: N } },
    select: { id: true },
  });
  const roleIds = roles.map((r) => r.id);
  if (roleIds.length > 0) {
    await prisma.role_permissions.deleteMany({
      where: { role_id: { in: roleIds } },
    });
    await prisma.roles.deleteMany({ where: { id: { in: roleIds } } });
  }
}

beforeAll(async () => {
  await cleanup();
  savedBudget = await getBudgetUsd();

  app = Fastify({ logger: false });
  await app.register(cookie);
  await app.register(rateLimit, { max: 1000, timeWindow: "1 minute" });
  app.addHook("onRequest", securityHeaders);
  await app.register(aiRoutes, { prefix: "/api/admin/ai" });

  const admin = await prisma.users.findFirst({
    where: { roles: { name: "admin" } },
  });
  if (!admin) throw new Error("Test setup: admin user not found");
  adminToken = token({
    id: admin.id,
    username: admin.username,
    roleName: "admin",
  });

  // Ordinary ai.use holder — may chat, must NOT control the budget.
  const role = await prisma.roles.create({
    data: { name: `${N}aiuse_${stamp}`, display_name: "AI Use Only" },
  });
  aiRoleId = role.id;
  const perm = await prisma.permissions.findFirst({
    where: { name: "ai.use" },
  });
  if (!perm) throw new Error("Test setup: ai.use permission missing");
  await prisma.role_permissions.create({
    data: { role_id: aiRoleId, permission_id: perm.id },
  });
  const user = await prisma.users.create({
    data: {
      username: `${N}user_${stamp}`,
      email: `${N}${stamp}@test.local`,
      password_hash:
        "$2a$10$invalidinvalidinvalidinvalidinvalidinvalidinvalidinvali",
      first_name: "AiUse",
      last_name: "Only",
      is_active: true,
      role_id: aiRoleId,
    },
  });
  aiUserToken = token({
    id: user.id,
    username: user.username,
    roleName: role.name,
  });
});

afterAll(async () => {
  await setBudgetUsd(savedBudget); // restore whatever the test DB had
  if (app) await app.close();
  await cleanup();
  await prisma.$disconnect();
});

describe("PUT /ai/budget permission gate (audit L8a)", () => {
  it("rejects an ordinary ai.use holder with 403", async () => {
    const res = await app.inject({
      method: "PUT",
      url: "/api/admin/ai/budget",
      headers: {
        Authorization: `Bearer ${aiUserToken}`,
        "Content-Type": "application/json",
      },
      payload: { budget_usd: 0 },
    });
    expect(res.statusCode).toBe(403);
    // ...and the budget must be untouched.
    expect(await getBudgetUsd()).toBe(savedBudget);
  });

  it("allows an admin to set the budget", async () => {
    const res = await app.inject({
      method: "PUT",
      url: "/api/admin/ai/budget",
      headers: {
        Authorization: `Bearer ${adminToken}`,
        "Content-Type": "application/json",
      },
      payload: { budget_usd: savedBudget },
    });
    expect(res.statusCode).toBe(200);
  });
});