NianAIGC/tests/usage-service.test.ts

import { mkdtemp, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import {
  createGenerationJob,
  deleteGenerationJob,
  listUsageEvents,
  recordUsageEvent,
  recordUsageForJob
} from "@/lib/server/data-store";
import { getAdminUsageReport, getPersonalUsageReport } from "@/lib/server/usage-service";
import { usagePresetRange } from "@/lib/usage";
import type { GenerationJob, UsageContext } from "@/lib/types";

let runtimeDir = "";
const previousEnv = new Map<string, string | undefined>();
const envNames = ["ZHINIAN_RUNTIME_DIR", "ZHINIAN_DATA_BACKEND"];

describe("usage metering and reports", () => {
  beforeEach(async () => {
    runtimeDir = await mkdtemp(join(tmpdir(), "zhinian-usage-"));
    for (const name of envNames) previousEnv.set(name, process.env[name]);
    process.env.ZHINIAN_RUNTIME_DIR = runtimeDir;
    process.env.ZHINIAN_DATA_BACKEND = "local";
  });

  afterEach(async () => {
    for (const name of envNames) restoreEnv(name, previousEnv.get(name));
    previousEnv.clear();
    await rm(runtimeDir, { force: true, recursive: true });
  });

  it("uses China Standard Time for preset boundaries", () => {
    const range = usagePresetRange("month", new Date("2026-07-27T16:30:00.000Z"));
    expect(range).toMatchObject({
      startDate: "2026-07-01",
      endDate: "2026-07-28",
      from: "2026-07-01T00:00:00+08:00",
      to: "2026-07-29T00:00:00+08:00",
      dayCount: 28
    });
  });

  it("records one immutable event per real platform job", async () => {
    const job = await createUsageJob("job-platform", platformContext("owner-a", "tenant-a", "org-a"));
    const first = await recordUsageForJob(job);
    const second = await recordUsageForJob(job);
    expect(second?.id).toBe(first?.id);
    expect(await listUsageEvents({ source: "platform" })).toHaveLength(1);

    await deleteGenerationJob(job.id);
    const retained = await listUsageEvents({ source: "platform" });
    expect(retained).toHaveLength(1);
    expect(retained[0]).toMatchObject({
      jobId: job.id,
      quantity: 1,
      estimatedUnit: "job",
      organizationId: "org-a"
    });
  });

  it("does not meter mock jobs or public API clients", async () => {
    const mock = await createUsageJob("job-mock", platformContext("owner-a"), { provider: "mock" });
    const api = await createUsageJob("job-api", undefined, { externalClientId: "partner-a" });
    expect(await recordUsageForJob(mock)).toBeNull();
    expect(await recordUsageForJob(api)).toBeNull();
    expect(await listUsageEvents()).toHaveLength(0);
  });

  it("builds personal and administrator reports from distinct successful jobs", async () => {
    await recordUsageEvent({
      ownerId: "owner-a",
      jobId: "job-a",
      source: "platform",
      capability: "image.generate",
      provider: "bailian",
      reqKey: "wan2.2-t2i-plus",
      accountDisplayName: "账号 A",
      tenantId: "tenant-a",
      organizationId: "org-a",
      organizationName: "组织 A",
      quantity: 4,
      estimatedUnit: "image",
      createdAt: "2026-07-03T01:00:00.000Z"
    });
    await recordUsageEvent({
      ownerId: "owner-b",
      jobId: "job-b",
      source: "platform",
      capability: "video.generate",
      provider: "seedance",
      accountDisplayName: "账号 B",
      quantity: 1,
      estimatedUnit: "job",
      createdAt: "2026-07-04T01:00:00.000Z"
    });
    await recordUsageEvent({
      ownerId: "api:partner-a",
      jobId: "job-api",
      source: "api",
      capability: "image.generate",
      provider: "bailian",
      quantity: 1,
      estimatedUnit: "job",
      createdAt: "2026-07-05T01:00:00.000Z"
    });
    await recordUsageEvent({
      ownerId: "owner-a",
      jobId: "job-mock",
      source: "platform",
      capability: "image.generate",
      provider: "mock",
      quantity: 1,
      estimatedUnit: "job",
      createdAt: "2026-07-06T01:00:00.000Z"
    });

    const now = new Date("2026-07-28T04:00:00.000Z");
    const personal = await getPersonalUsageReport("owner-a", "month", now);
    expect(personal.total).toBe(1);
    expect(personal.byCapability.find((item) => item.key === "image.generate")?.count).toBe(1);

    const report = await getAdminUsageReport({}, [{
      organizationId: "org-a",
      organizationName: "组织 A",
      organizationBindTenantId: 1
    }], now);
    expect(report.summary).toMatchObject({ total: 2, activeAccounts: 2, activeOrganizations: 1 });
    expect(report.organizations.map((row) => row.organizationName)).toEqual(["组织 A", "未归属组织"]);

    const filtered = await getAdminUsageReport({ organizationId: "org-a" }, [], now);
    expect(filtered.summary.total).toBe(1);
    expect(filtered.accounts[0]?.accountName).toBe("账号 A");
    expect(filtered.options.accounts.map((option) => option.value)).toEqual(["owner-a"]);
  });
});

async function createUsageJob(
  id: string,
  usageContext?: UsageContext,
  overrides: Partial<GenerationJob> = {}
): Promise<GenerationJob> {
  return createGenerationJob({
    id,
    ownerId: usageContext?.accountId || "api:partner-a",
    externalClientId: overrides.externalClientId,
    capability: "image.generate",
    provider: overrides.provider || "bailian",
    reqKey: "wan2.2-t2i-plus",
    status: "running",
    inputAssetIds: [],
    inputUrls: [],
    outputAssetIds: [],
    requestPayload: {},
    usageContext
  });
}

function platformContext(accountId: string, tenantId?: string, organizationId?: string): UsageContext {
  return {
    source: "platform",
    accountId,
    username: accountId,
    displayName: accountId,
    tenantId,
    organizationId,
    organizationName: organizationId ? "组织 A" : undefined
  };
}

function restoreEnv(name: string, value: string | undefined) {
  if (value === undefined) {
    delete process.env[name];
    return;
  }
  process.env[name] = value;
}