aegida-console / lib / chat / quota.test.ts
quota.test.ts
Raw
// @vitest-environment node

import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";

vi.mock("server-only", () => ({}));

import {
  getGateQuota,
  parseGateQuota,
  QuotaError,
} from "@/lib/chat/quota";

const finiteQuota = {
  object: "aegida.quota",
  model: "qwen",
  available_tokens: 42_000,
  unlimited: false,
  exhausted: false,
  resets_at: "2026-08-05T00:00:00Z",
  as_of: "2026-08-04T12:00:00Z",
};

beforeEach(() => {
  vi.stubEnv("AEGIDA_GATE_URL", "https://gate.example.test");
});

afterEach(() => {
  vi.unstubAllEnvs();
  vi.unstubAllGlobals();
});

describe("Aegida effective quota", () => {
  it("strictly parses finite, exhausted, and unlimited snapshots", () => {
    expect(parseGateQuota(finiteQuota, "qwen")).toEqual({
      object: "aegida.quota",
      model: "qwen",
      availableTokens: 42_000,
      unlimited: false,
      exhausted: false,
      resetsAt: "2026-08-05T00:00:00Z",
      asOf: "2026-08-04T12:00:00Z",
    });
    expect(
      parseGateQuota(
        { ...finiteQuota, available_tokens: 0, exhausted: true, resets_at: null },
        "qwen",
      ),
    ).toMatchObject({ availableTokens: 0, exhausted: true, unlimited: false });
    expect(
      parseGateQuota(
        {
          ...finiteQuota,
          model: "auto",
          available_tokens: null,
          unlimited: true,
          exhausted: false,
          resets_at: null,
        },
        "auto",
      ),
    ).toMatchObject({ availableTokens: null, unlimited: true, exhausted: false });
  });

  it.each([
    ["unknown field", { ...finiteQuota, private_policy: "sentinel" }],
    ["wrong object", { ...finiteQuota, object: "quota" }],
    ["wrong model", finiteQuota, "auto"],
    ["negative tokens", { ...finiteQuota, available_tokens: -1 }],
    ["zero not exhausted", { ...finiteQuota, available_tokens: 0 }],
    ["positive exhausted", { ...finiteQuota, exhausted: true }],
    ["finite null tokens", { ...finiteQuota, available_tokens: null }],
    ["unlimited tokens", { ...finiteQuota, unlimited: true }],
    [
      "unlimited reset",
      { ...finiteQuota, available_tokens: null, unlimited: true, resets_at: "2026-08-05T00:00:00Z" },
    ],
    ["bad reset", { ...finiteQuota, resets_at: "tomorrow" }],
    ["bad as-of", { ...finiteQuota, as_of: "today" }],
  ])("rejects %s", (_name, value, model = "qwen") => {
    expect(() => parseGateQuota(value, model)).toThrow(QuotaError);
  });

  it("relays the exact user token with model and no-store semantics", async () => {
    const fetchMock = vi.fn().mockResolvedValue(Response.json(finiteQuota));
    vi.stubGlobal("fetch", fetchMock);
    const signal = new AbortController().signal;

    await expect(getGateQuota("qwen", "Bearer exact-user-jwt", signal)).resolves.toMatchObject({
      model: "qwen",
      availableTokens: 42_000,
    });
    expect(fetchMock).toHaveBeenCalledWith(
      "https://gate.example.test/v1/aegida/quota?model=qwen",
      {
        method: "GET",
        headers: { Accept: "application/json", Authorization: "Bearer exact-user-jwt" },
        cache: "no-store",
        redirect: "error",
        signal,
      },
    );
  });

  it("classifies upstream authentication and sanitizes every other failure", async () => {
    vi.stubGlobal(
      "fetch",
      vi
        .fn()
        .mockResolvedValueOnce(new Response("private-auth-sentinel", { status: 401 }))
        .mockResolvedValueOnce(new Response("private-dependency-sentinel", { status: 500 })),
    );

    await expect(
      getGateQuota("qwen", "Bearer user", new AbortController().signal),
    ).rejects.toMatchObject({ reason: "unauthorized" });
    let error: unknown;
    try {
      await getGateQuota("qwen", "Bearer user", new AbortController().signal);
    } catch (caught) {
      error = caught;
    }
    expect(error).toMatchObject({ reason: "unavailable" });
    expect(String(error)).not.toContain("private-dependency-sentinel");
  });
});