1
import { afterEach, describe, expect, it, vi } from "vitest";
2
 
3
import { BranchName, ProjectId, RunId } from "@/contracts";
4
import type { RunMetaState } from "@/worker/contracts";
5
import {
6
  ensureRunCancelRequested,
7
  ensureRunCanceling,
8
  markOwnershipLost,
9
} from "@/worker/dispatch/shared/run-execution-context/control";
10
import type { RunStore } from "@/worker/dispatch/shared/run-execution-context/types";
11
 
12
import { createQueueScope, createQueueState } from "../../../helpers/dispatch/shared";
13
 
14
const makeScope = () => createQueueScope();
15
 
16
const makeState = () => createQueueState();
17
 
18
const makeRunMeta = (status: RunMetaState["status"], startedAt: RunMetaState["startedAt"] = null): RunMetaState => ({
19
  runId: RunId.assertDecode("run_0000000000000000000000"),
20
  projectId: ProjectId.assertDecode("prj_0000000000000000000000"),
21
  status,
22
  triggerType: "manual",
23
  branch: BranchName.assertDecode("main"),
24
  commitSha: null,
25
  currentStep: null,
26
  startedAt,
27
  finishedAt: null,
28
  exitCode: null,
29
  errorMessage: null,
30
});
31
 
32
afterEach(() => {
33
  vi.restoreAllMocks();
34
});
35
 
36
describe("queue control helpers", () => {
37
  it("moves a running run to cancel_requested and backfills startedAt", async () => {
38
    vi.spyOn(console, "log").mockImplementation(() => {});
39
    const scope = makeScope();
40
    const state = makeState();
41
    const current = makeRunMeta("running");
42
    const runStore = {
43
      getMeta: vi.fn(async () => current),
44
      tryUpdateState: vi.fn(async (input) => ({
45
        kind: "applied" as const,
46
        state: {
47
          ...current,
48
          ...input,
49
          runId: current.runId,
50
          projectId: current.projectId,
51
        },
52
      })),
53
    } as Pick<RunStore, "getMeta" | "tryUpdateState"> as RunStore;
54
 
55
    const result = await ensureRunCancelRequested(scope, state, runStore);
56
 
57
    expect(runStore.tryUpdateState).toHaveBeenCalledWith(
58
      expect.objectContaining({
59
        status: "cancel_requested",
60
        startedAt: scope.startedAt,
61
      }),
62
    );
63
    expect(result?.status).toBe("cancel_requested");
64
    expect(result?.startedAt).toBe(scope.startedAt);
65
  });
66
 
67
  it("moves cancel_requested to canceling without changing the rest of the transition data", async () => {
68
    vi.spyOn(console, "log").mockImplementation(() => {});
69
    const scope = makeScope();
70
    const state = makeState();
71
    const current = makeRunMeta("cancel_requested", scope.startedAt);
72
    const runStore = {
73
      getMeta: vi.fn(async () => current),
74
      tryUpdateState: vi.fn(async (input) => ({
75
        kind: "applied" as const,
76
        state: {
77
          ...current,
78
          ...input,
79
          runId: current.runId,
80
          projectId: current.projectId,
81
        },
82
      })),
83
    } as Pick<RunStore, "getMeta" | "tryUpdateState"> as RunStore;
84
 
85
    const result = await ensureRunCanceling(scope, state, runStore);
86
 
87
    expect(runStore.tryUpdateState).toHaveBeenCalledWith(
88
      expect.objectContaining({
89
        status: "canceling",
90
        startedAt: scope.startedAt,
91
      }),
92
    );
93
    expect(result?.status).toBe("canceling");
94
  });
95
 
96
  it("marks ownership loss only once", () => {
97
    vi.spyOn(console, "log").mockImplementation(() => {});
98
    const scope = makeScope();
99
    const state = makeState();
100
 
101
    markOwnershipLost(scope, state, "failed");
102
    markOwnershipLost(scope, state, "active");
103
 
104
    expect(state.ownershipLost).toBe(true);
105
    expect(state.ownershipLossStatus).toBe("failed");
106
  });
107
});