UHD_Media_Tool/packages/tasks/test/op-ledger.test.ts

import { sql } from "kysely";
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { FencingError, getOperation, migrateTasksTables, withAttemptToken } from "../src";
import { connectDb, dropSchemas, pgAvailable, resetSchemas, type TestDb, testSchemas } from "./helpers";

const schemas = testSchemas("ledger");
const hasPg = await pgAvailable();

describe.skipIf(!hasPg)("withAttemptToken(操作账本 fencing)", () => {
  let db: TestDb;

  beforeAll(async () => {
    await resetSchemas(schemas);
    db = connectDb(schemas.biz);
    await migrateTasksTables(db);
    await sql`CREATE TABLE download_task (id text PRIMARY KEY, hash text)`.execute(db);
  });

  afterAll(async () => {
    await db?.destroy();
    await dropSchemas(schemas);
  });

  it("成功后记录结果;再次调用直接返回结果不重做", async () => {
    let calls = 0;
    const run = () =>
      withAttemptToken(
        db,
        "op-1",
        async (info) => {
          calls++;
          return { hash: "h1", attempt: info.attempt };
        },
        {
          runId: "r1",
          commit: async (trx, result) => {
            await sql`INSERT INTO download_task (id, hash) VALUES ('d1', ${result.hash})`.execute(trx);
          },
        },
      );
    expect(await run()).toEqual({ hash: "h1", attempt: 1 });
    expect(await run()).toEqual({ hash: "h1", attempt: 1 });
    expect(calls).toBe(1);
    const row = await getOperation(db, "op-1");
    expect(row).toMatchObject({ status: "succeeded", attempts: 1, runId: "r1" });
  });

  it("租约丢失:旧 attempt 执行完后提交被拒绝,新 attempt 的结果生效", async () => {
    let releaseOld!: () => void;
    const oldGate = new Promise<void>((r) => {
      releaseOld = r;
    });
    let oldStarted!: () => void;
    const oldStartedP = new Promise<void>((r) => {
      oldStarted = r;
    });

    // 旧 worker(被 SIGSTOP 之类卡住):已开始执行,尚未提交
    const oldAttempt = withAttemptToken(
      db,
      "op-2",
      async () => {
        oldStarted();
        await oldGate;
        return "old";
      },
      {
        commit: async (trx) => void (await sql`INSERT INTO download_task (id, hash) VALUES ('d2', 'old')`.execute(trx)),
      },
    );
    await oldStartedP;

    // 新 worker 接管同一操作:resumed=true,需要先对账外部状态
    let resumed = false;
    const fresh = await withAttemptToken(
      db,
      "op-2",
      async (info) => {
        resumed = info.resumed;
        return "new";
      },
      {
        commit: async (trx) => void (await sql`INSERT INTO download_task (id, hash) VALUES ('d2', 'new')`.execute(trx)),
      },
    );
    expect(fresh).toBe("new");
    expect(resumed).toBe(true);

    releaseOld();
    await expect(oldAttempt).rejects.toBeInstanceOf(FencingError);
    const r = await sql<{ hash: string }>`SELECT hash FROM download_task WHERE id = 'd2'`.execute(db);
    expect(r.rows).toEqual([{ hash: "new" }]);
    expect(await getOperation(db, "op-2")).toMatchObject({ status: "succeeded", attempts: 2, result: "new" });
  });

  it("fn 抛错:记录 last_error,不写结果,可重试", async () => {
    await expect(
      withAttemptToken(db, "op-3", async () => {
        throw new Error("qb timeout");
      }),
    ).rejects.toThrow("qb timeout");
    expect(await getOperation(db, "op-3")).toMatchObject({ status: "started", lastError: "qb timeout" });
    expect(await withAttemptToken(db, "op-3", async (i) => i.attempt)).toBe(2);
  });
});