Files
NianAIGC/tests/data-store-concurrency.test.ts

65 lines
2.1 KiB
TypeScript

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,
listGenerationJobs,
updateGenerationJob
} from "@/lib/server/data-store";
import { DEFAULT_OWNER_ID } from "@/lib/server/runtime";
let runtimeDir = "";
let previousRuntimeDir: string | undefined;
let previousDataBackend: string | undefined;
describe("local data store concurrency", () => {
beforeEach(async () => {
runtimeDir = await mkdtemp(join(tmpdir(), "zhinian-store-"));
previousRuntimeDir = process.env.ZHINIAN_RUNTIME_DIR;
previousDataBackend = process.env.ZHINIAN_DATA_BACKEND;
process.env.ZHINIAN_RUNTIME_DIR = runtimeDir;
process.env.ZHINIAN_DATA_BACKEND = "local";
});
afterEach(async () => {
restoreEnv("ZHINIAN_RUNTIME_DIR", previousRuntimeDir);
restoreEnv("ZHINIAN_DATA_BACKEND", previousDataBackend);
await rm(runtimeDir, { force: true, recursive: true });
});
it("serializes concurrent job writes without losing records", async () => {
const created = await Promise.all(Array.from({ length: 8 }, (_, 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 }
})));
await Promise.all(created.map((job, index) => updateGenerationJob(job.id, {
status: "succeeded",
outputAssetIds: [`asset-${index}`]
})));
const stored = await listGenerationJobs(DEFAULT_OWNER_ID, 20);
const storedById = new Map(stored.map((job) => [job.id, job]));
for (const job of created) {
expect(storedById.get(job.id)?.status).toBe("succeeded");
}
});
});
function restoreEnv(name: string, value: string | undefined) {
if (value === undefined) {
delete process.env[name];
return;
}
process.env[name] = value;
}