9router/tests/unit/user-token-limit-enforcer.test.js

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

const getUserById = vi.fn();
const getUserProviderTokenUsageSince = vi.fn();
const getUserProviderEarliestTokenUsageSince = vi.fn();
const getUserTokenQuotaSession = vi.fn();
const ensureUserTokenQuotaSession = vi.fn();
const getUserTokenLimits = vi.fn();

vi.mock("@/lib/db/index.js", () => ({
  getUserById,
  ensureUserTokenQuotaSession,
  getUserProviderEarliestTokenUsageSince,
  getUserProviderTokenUsageSince,
  getUserTokenQuotaSession,
  getUserTokenLimits,
}));

const {
  checkUserTokenLimit,
  getUserTokenLimitWindowStart,
} = await import("@/lib/tokenLimitEnforcer.js");

describe("user token limit enforcement", () => {
  beforeEach(() => {
    getUserById.mockReset();
    ensureUserTokenQuotaSession.mockReset();
    getUserProviderEarliestTokenUsageSince.mockReset();
    getUserProviderTokenUsageSince.mockReset();
    getUserTokenQuotaSession.mockReset();
    getUserTokenLimits.mockReset();
    getUserById.mockResolvedValue({ id: "user-1", role: "user", isActive: true });
    getUserTokenLimits.mockResolvedValue({
      "orbit-provider": { session: 100, weekly: 1000 },
      codex: { session: 200, weekly: 2000 },
    });
    getUserTokenQuotaSession.mockResolvedValue("2026-07-17T06:00:00.000Z");
  });

  it("calculates rolling session and Monday Vietnam weekly window starts", () => {
    const now = new Date("2026-07-17T10:00:00.000Z");

    expect(getUserTokenLimitWindowStart("session", now).toISOString())
      .toBe("2026-07-17T05:00:00.000Z");
    expect(getUserTokenLimitWindowStart("weekly", now).toISOString())
      .toBe("2026-07-12T17:00:00.000Z");
  });

  it("blocks when the rolling session total reaches its limit", async () => {
    getUserProviderTokenUsageSince.mockResolvedValueOnce(100);

    const result = await checkUserTokenLimit(
      "user-1",
      "orbit-provider",
      new Date("2026-07-17T10:00:00.000Z"),
    );

    expect(result).toMatchObject({
      exceeded: true,
      provider: "orbit-provider",
      windowType: "session",
      limit: 100,
      used: 100,
    });
    expect(getUserProviderTokenUsageSince).toHaveBeenCalledTimes(1);
  });

  it("does not block from an expired fixed session until a new request begins one", async () => {
    getUserTokenQuotaSession.mockResolvedValue("2026-07-17T05:00:00.000Z");
    getUserProviderTokenUsageSince.mockResolvedValueOnce(0);

    await expect(checkUserTokenLimit(
      "user-1",
      "orbit-provider",
      new Date("2026-07-17T10:10:00.000Z"),
    )).resolves.toBeNull();

    expect(getUserProviderTokenUsageSince).toHaveBeenCalledWith(
      "user-1",
      "orbit-provider",
      new Date("2026-07-12T17:00:00.000Z"),
    );
  });

  it("checks weekly usage after the session window still has headroom", async () => {
    getUserProviderTokenUsageSince
      .mockResolvedValueOnce(80)
      .mockResolvedValueOnce(1200);

    const result = await checkUserTokenLimit(
      "user-1",
      "orbit-provider",
      new Date("2026-07-17T10:00:00.000Z"),
    );

    expect(result).toMatchObject({ windowType: "weekly", limit: 1000, used: 1200 });
    expect(getUserProviderTokenUsageSince).toHaveBeenCalledTimes(2);
  });

  it("exempts administrators, unknown users, and unsupported providers", async () => {
    getUserById.mockResolvedValueOnce({ id: "admin-1", role: "admin", isActive: true });
    await expect(checkUserTokenLimit("admin-1", "codex")).resolves.toBeNull();

    getUserById.mockResolvedValueOnce(null);
    await expect(checkUserTokenLimit("missing", "codex")).resolves.toBeNull();

    await expect(checkUserTokenLimit("user-1", "openai")).resolves.toBeNull();
    expect(getUserTokenLimits).not.toHaveBeenCalled();
  });

  it("treats zero limits as unlimited", async () => {
    getUserTokenLimits.mockResolvedValue({
      "orbit-provider": { session: 0, weekly: 0 },
      codex: { session: 0, weekly: 0 },
    });

    await expect(checkUserTokenLimit("user-1", "codex")).resolves.toBeNull();
    expect(getUserProviderTokenUsageSince).not.toHaveBeenCalled();
  });
});