1
import { vi } from "vitest";
2
 
3
import { type CommitSha as CommitShaType, ProjectId, RunId, UnixTimestampMs } from "@/contracts";
4
import { type ProjectRunStatus, type RecordRunResolvedCommitResult, expectTrusted } from "@/worker/contracts";
5
import type { RunExecutionContext } from "@/worker/dispatch/shared/run-execution-context";
6
 
7
import {
8
  createQueueProjectControlStub,
9
  createQueueRunControlStub,
10
  createQueueRunRuntimeStub,
11
  createQueueRunStoreStub,
12
  createQueueScope,
13
  createQueueState,
14
  createQueueLeaseStub,
15
} from "./dispatch/shared";
16
 
17
type RunEnvironmentHarnessContext = Pick<
18
  RunExecutionContext,
19
  "control" | "projectControl" | "runStore" | "runtime" | "scope" | "state"
20
>;
21
 
22
export type ExecResult = {
23
  success: boolean;
24
  stdout: string;
25
  stderr: string;
26
};
27
 
28
export const INVALID_REPO_CONFIG = "run:\n  steps: [";
29
export const VALID_REPO_CONFIG = `version: 1
30
checkout:
31
  depth: 3
32
run:
33
  workingDirectory: .
34
  timeoutSeconds: 60
35
  steps:
36
    - name: test
37
      run: echo hello
38
`;
39
 
40
export const createLeaseStub = createQueueLeaseStub;
41
 
42
export const createExecResult = (overrides: Partial<ExecResult> = {}): ExecResult => ({
43
  success: true,
44
  stdout: "",
45
  stderr: "",
46
  ...overrides,
47
});
48
 
49
export const createHarness = ({
50
  projectIdValue,
51
  runIdValue,
52
  snapshotCommitSha,
53
  runMetaCommitSha = snapshotCommitSha,
54
  recordRunResolvedCommitResult = { kind: "applied" },
55
  configContent = INVALID_REPO_CONFIG,
56
  execImpl,
57
}: {
58
  projectIdValue: string;
59
  runIdValue: string;
60
  snapshotCommitSha: CommitShaType | null;
61
  runMetaCommitSha?: CommitShaType | null;
62
  recordRunResolvedCommitResult?: RecordRunResolvedCommitResult;
63
  configContent?: string;
64
  execImpl: (command: string) => Promise<ExecResult> | ExecResult;
65
}) => {
66
  const projectId = ProjectId.assertDecode(projectIdValue);
67
  const runId = RunId.assertDecode(runIdValue);
68
  const events: string[] = [];
69
 
70
  const checkoutSession = {
71
    id: "checkout-session",
72
    exec: vi.fn(async (command: string) => execImpl(command)),
73
    readFile: vi.fn(async () => {
74
      events.push("read-file");
75
      return {
76
        content: configContent,
77
      };
78
    }),
79
  };
80
  const runSession = {
81
    id: "run-session",
82
  };
83
 
84
  const recordRunResolvedCommit = vi.fn(
85
    async (_input: { projectId: ProjectId; runId: RunId; commitSha: CommitShaType }) => {
86
      events.push("record-run-resolved-commit");
87
      return recordRunResolvedCommitResult;
88
    },
89
  );
90
  const replaceSteps = vi.fn(async () => {});
91
  let sessionCreations = 0;
92
 
93
  const sandbox = {
94
    setKeepAlive: vi.fn(async () => {}),
95
    createSession: vi.fn(async () => {
96
      sessionCreations += 1;
97
      return sessionCreations === 1 ? checkoutSession : runSession;
98
    }),
99
    deleteSession: vi.fn(async () => {
100
      events.push("delete-session");
101
    }),
102
  };
103
 
104
  const startedAt = expectTrusted(UnixTimestampMs, Date.now(), "UnixTimestampMs");
105
  const scope = createQueueScope({
106
    projectId,
107
    runId,
108
    snapshotCommitSha,
109
    startedAt,
110
  });
111
  const state = createQueueState();
112
 
113
  const context: RunEnvironmentHarnessContext = {
114
    scope,
115
    state,
116
    projectControl: createQueueProjectControlStub({
117
      recordHeartbeat: vi.fn(async () => null),
118
      recordResolvedCommit: async (commitSha: CommitShaType) =>
119
        await recordRunResolvedCommit({
120
          projectId,
121
          runId,
122
          commitSha,
123
        }),
124
    }),
125
    runStore: createQueueRunStoreStub({
126
      getMeta: vi.fn(async () => ({
127
        runId,
128
        projectId,
129
        status: "queued" as const,
130
        triggerType: "manual" as const,
131
        branch: scope.snapshot.branch,
132
        commitSha: runMetaCommitSha,
133
        currentStep: null,
134
        startedAt: null,
135
        finishedAt: null,
136
        exitCode: null,
137
        errorMessage: null,
138
      })),
139
      replaceSteps,
140
    }),
141
    runtime: createQueueRunRuntimeStub({
142
      sandbox: sandbox as unknown as RunEnvironmentHarnessContext["runtime"]["sandbox"],
143
    }),
144
    control: createQueueRunControlStub({
145
      markOwnershipLost: (status: ProjectRunStatus | null) => {
146
        state.ownershipLost = true;
147
        state.ownershipLossStatus = status;
148
        events.push("ownership-lost");
149
      },
150
    }),
151
  };
152
 
153
  return {
154
    context,
155
    checkoutSession,
156
    events,
157
    projectId,
158
    recordRunResolvedCommit,
159
    replaceSteps,
160
    runId,
161
    runSession,
162
    sandbox,
163
  };
164
};