NianAIGC/tests/task-management.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 {
  authenticatePublicApiRequest,
  PublicApiAuthError,
  type PublicApiClient
} from "@/lib/server/public-api-auth";
import { createPublicGenerationJob, PublicApiConflictError } from "@/lib/server/public-api-jobs";
import { claimGenerationJobs, createGenerationJob, getGenerationJob, listGenerationJobs } from "@/lib/server/data-store";
import { runWorkerTick } from "@/lib/server/task-manager";
import { DEFAULT_OWNER_ID } from "@/lib/server/runtime";
import { signWebhookBody } from "@/lib/server/webhook";

let runtimeDir = "";
const previousEnv = new Map<string, string | undefined>();
const envNames = [
  "ZHINIAN_RUNTIME_DIR",
  "NEXT_PUBLIC_SUPABASE_URL",
  "SUPABASE_SERVICE_ROLE_KEY",
  "ZHINIAN_API_KEYS",
  "JIMENG_VISUAL_MOCK",
  "VOLCENGINE_ACCESS_KEY_ID",
  "VOLCENGINE_SECRET_ACCESS_KEY",
  "ZHINIAN_WEBHOOK_SECRET"
];

describe("task management and public API helpers", () => {
  beforeEach(async () => {
    runtimeDir = await mkdtemp(join(tmpdir(), "zhinian-tasks-"));
    for (const name of envNames) previousEnv.set(name, process.env[name]);
    process.env.ZHINIAN_RUNTIME_DIR = runtimeDir;
    process.env.ZHINIAN_API_KEYS = "agent-a:secret-a,agent-b:secret-b";
    process.env.JIMENG_VISUAL_MOCK = "true";
    delete process.env.NEXT_PUBLIC_SUPABASE_URL;
    delete process.env.SUPABASE_SERVICE_ROLE_KEY;
    delete process.env.VOLCENGINE_ACCESS_KEY_ID;
    delete process.env.VOLCENGINE_SECRET_ACCESS_KEY;
  });

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

  it("authenticates public API requests with bearer and header keys", () => {
    expect(authenticatePublicApiRequest(new Request("http://local.test", {
      headers: { Authorization: "Bearer secret-a" }
    })).id).toBe("agent-a");
    expect(authenticatePublicApiRequest(new Request("http://local.test", {
      headers: { "X-Zhinian-Api-Key": "secret-b" }
    })).id).toBe("agent-b");
    expect(() => authenticatePublicApiRequest(new Request("http://local.test"))).toThrow(PublicApiAuthError);
  });

  it("deduplicates public job creation by idempotency key and rejects conflicts", async () => {
    const client: PublicApiClient = { id: "agent-a", key: "secret-a" };
    const request = new Request("http://local.test/api/v1/jobs", {
      headers: { "Idempotency-Key": "idem-1" }
    });
    const first = await createPublicGenerationJob({
      client,
      request,
      origin: "http://local.test",
      body: {
        capability: "image.generate",
        prompt: "生成一张专业产品主图"
      }
    });
    const second = await createPublicGenerationJob({
      client,
      request,
      origin: "http://local.test",
      body: {
        capability: "image.generate",
        prompt: "生成一张专业产品主图"
      }
    });
    expect(second.reused).toBe(true);
    expect(second.job.id).toBe(first.job.id);
    await expect(createPublicGenerationJob({
      client,
      request,
      origin: "http://local.test",
      body: {
        capability: "image.generate",
        prompt: "不同的提示词"
      }
    })).rejects.toBeInstanceOf(PublicApiConflictError);
  });

  it("claims local jobs without duplicate ownership", async () => {
    await Promise.all(Array.from({ length: 6 }, (_, index) => createGenerationJob({
      ownerId: DEFAULT_OWNER_ID,
      capability: "image.generate",
      provider: "mock",
      reqKey: "jimeng_seedream46_cvtob",
      status: "queued",
      prompt: `job ${index}`,
      inputAssetIds: [],
      inputUrls: [],
      outputAssetIds: [],
      requestPayload: { index },
      priority: index
    })));

    const [left, right] = await Promise.all([
      claimGenerationJobs({ workerId: "worker-a", limit: 3 }),
      claimGenerationJobs({ workerId: "worker-b", limit: 3 })
    ]);
    const claimedIds = [...left, ...right].map((job) => job.id);
    expect(new Set(claimedIds).size).toBe(6);
    expect(await claimGenerationJobs({ workerId: "worker-c", limit: 1 })).toHaveLength(0);
  });

  it("processes a mock queued job to a terminal result through the worker tick", async () => {
    const job = await createPublicGenerationJob({
      client: { id: "agent-a", key: "secret-a" },
      request: new Request("http://local.test/api/v1/jobs"),
      origin: "http://local.test",
      body: {
        capability: "image.generate",
        prompt: "生成一张适合社媒传播的品牌主视觉"
      }
    });
    expect(job.job.status).toBe("queued");
    const tick = await runWorkerTick({
      workerId: "test-worker",
      origin: "http://local.test",
      limit: 1
    });
    expect(tick.claimed).toBe(1);
    const stored = await getGenerationJob(job.job.id);
    expect(stored?.status).toBe("succeeded");
    expect(stored?.completedAt).toBeTruthy();
    expect(stored?.outputAssetIds.length).toBe(1);
    expect(await listGenerationJobs(DEFAULT_OWNER_ID, 10)).toHaveLength(1);
  });

  it("signs webhook bodies with the configured secret", () => {
    expect(signWebhookBody("{\"ok\":true}", "secret")).toBe(
      "sha256=f6b4a2841c93f8bf2fb8f2c13d8fb0b6c8e8019f09ee405d248daa8385fad638"
    );
  });
});

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