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(); const envNames = [ "ZHINIAN_RUNTIME_DIR", "NEXT_PUBLIC_SUPABASE_URL", "SUPABASE_SERVICE_ROLE_KEY", "ZHINIAN_API_KEYS", "JIMENG_VISUAL_MOCK", "IMAGE_GENERATE_ENGINE", "EVOLINK_MOCK", "EVOLINK_API_KEY", "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.IMAGE_GENERATE_ENGINE; delete process.env.EVOLINK_MOCK; delete process.env.EVOLINK_API_KEY; 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("passes EvoLink quality through public image jobs", async () => { process.env.IMAGE_GENERATE_ENGINE = "evolink"; process.env.EVOLINK_MOCK = "true"; process.env.EVOLINK_API_KEY = "test-key"; const result = 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: "商品海报", quality: "high" } }); expect(result.job.provider).toBe("mock"); expect(result.job.requestPayload.providerPayload).toMatchObject({ model: "gpt-image-2", quality: "high" }); }); 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; }