1
import { runInDurableObject } from "cloudflare:test";
2
import { env } from "cloudflare:workers";
3
 
4
import { BranchName, DispatchMode, ExecutionRuntime } from "@/contracts";
5
import { ProjectDO } from "@/worker/durable";
6
import {
7
  acceptManualRun as acceptManualRunCommand,
8
  claimRunWork as claimRunWorkCommand,
9
  finalizeRunExecution as finalizeRunExecutionCommand,
10
  requestRunCancel as requestRunCancelCommand,
11
} from "@/worker/durable/project-do/commands";
12
import { getProjectConfigRow } from "@/worker/durable/project-do/repo";
13
import { transitionAcceptManualRun } from "@/worker/durable/project-do/transitions";
14
import { expectTrusted } from "@/worker/contracts";
15
import type { ProjectDoContext } from "@/worker/durable/project-do/types";
16
import { createLogger } from "@/worker/services/logger";
17
 
18
type ProjectDoInternals = Pick<ProjectDoContext, "ctx" | "env" | "db">;
19
 
20
export const getProjectDoInternals = (instance: ProjectDO): ProjectDoInternals =>
21
  instance as unknown as ProjectDoInternals;
22
 
23
export const createTestProjectDoContext = (instance: ProjectDO, envOverride?: Env): ProjectDoContext => {
24
  const { ctx, env, db } = getProjectDoInternals(instance);
25
 
26
  return {
27
    ctx,
28
    env: envOverride ?? env,
29
    db,
30
    logger: createLogger("test.project-do"),
31
    cacheProjectId: (projectId) => {
32
      (instance as unknown as { cachedProjectId?: string | null }).cachedProjectId = projectId;
33
    },
34
  };
35
};
36
 
37
export const createBoundStorageProxy = (
38
  storage: DurableObjectStorage,
39
  overrides: Partial<Pick<DurableObjectStorage, "getAlarm" | "setAlarm" | "transaction">>,
40
): DurableObjectStorage =>
41
  new Proxy(storage, {
42
    get(target, prop, receiver) {
43
      if (Object.prototype.hasOwnProperty.call(overrides, prop)) {
44
        return overrides[prop as keyof typeof overrides];
45
      }
46
 
47
      const value = Reflect.get(target, prop, receiver);
48
      return typeof value === "function" ? value.bind(target) : value;
49
    },
50
  });
51
 
52
export const createBoundTransactionProxy = (
53
  txn: DurableObjectTransaction,
54
  overrides: Partial<Pick<DurableObjectTransaction, "getAlarm" | "setAlarm">>,
55
): DurableObjectTransaction =>
56
  new Proxy(txn, {
57
    get(target, prop, receiver) {
58
      if (Object.prototype.hasOwnProperty.call(overrides, prop)) {
59
        return overrides[prop as keyof typeof overrides];
60
      }
61
 
62
      const value = Reflect.get(target, prop, receiver);
63
      return typeof value === "function" ? value.bind(target) : value;
64
    },
65
  });
66
 
67
export const withPatchedStorage = (context: ProjectDoContext, storage: DurableObjectStorage): ProjectDoContext => ({
68
  ...context,
69
  ctx: Object.assign(Object.create(context.ctx), { storage }) as DurableObjectState,
70
});
71
 
72
export const createAlarmSchedulingSuppressedContext = (context: ProjectDoContext): ProjectDoContext => {
73
  const transaction: DurableObjectStorage["transaction"] = async (closure) =>
74
    context.ctx.storage.transaction(
75
      async (txn) => await closure(createBoundTransactionProxy(txn, { setAlarm: async () => {} })),
76
    );
77
  const storage = createBoundStorageProxy(context.ctx.storage, { transaction });
78
 
79
  return withPatchedStorage(context, storage);
80
};
81
 
82
export const createAlarmWriteFailingContext = (
83
  context: ProjectDoContext,
84
  errorMessage = "alarm write failed",
85
): ProjectDoContext => {
86
  const transaction: DurableObjectStorage["transaction"] = async (closure) =>
87
    context.ctx.storage.transaction(async (txn) => {
88
      const failingTxn = createBoundTransactionProxy(txn, {
89
        setAlarm: async () => {
90
          throw new Error(errorMessage);
91
        },
92
      });
93
 
94
      return await closure(failingTxn);
95
    });
96
  const storage = createBoundStorageProxy(context.ctx.storage, { transaction });
97
 
98
  return withPatchedStorage(context, storage);
99
};
100
 
101
export const expectAcceptedManualRun = (
102
  result: Awaited<ReturnType<typeof acceptManualRunCommand>>,
103
): Extract<Awaited<ReturnType<typeof acceptManualRunCommand>>, { kind: "accepted" }> => {
104
  if (result.kind !== "accepted") {
105
    throw new Error(`Expected accepted manual run, got ${result.kind}:${result.reason}.`);
106
  }
107
 
108
  return result;
109
};
110
 
111
export const acceptManualRunWithoutAlarm = async (
112
  projectStub: ReturnType<typeof env.PROJECT_DO.getByName>,
113
  input: Parameters<typeof acceptManualRunCommand>[1],
114
): Promise<Extract<Awaited<ReturnType<typeof acceptManualRunCommand>>, { kind: "accepted" }>> => {
115
  let accepted: Awaited<ReturnType<typeof acceptManualRunCommand>> | null = null;
116
 
117
  await runInDurableObject(projectStub, async (instance: ProjectDO) => {
118
    const context = createAlarmSchedulingSuppressedContext(createTestProjectDoContext(instance));
119
    accepted = await acceptManualRunCommand(context, input);
120
  });
121
 
122
  if (!accepted) {
123
    throw new Error("acceptManualRunWithoutAlarm did not return an accepted run.");
124
  }
125
 
126
  return expectAcceptedManualRun(accepted);
127
};
128
 
129
export const fillManualRunQueueWithoutAlarmOrRunInitialization = async (
130
  projectStub: ReturnType<typeof env.PROJECT_DO.getByName>,
131
  input: Parameters<typeof acceptManualRunCommand>[1],
132
  count: number,
133
): Promise<void> => {
134
  await runInDurableObject(projectStub, async (instance: ProjectDO) => {
135
    const context = createAlarmSchedulingSuppressedContext(createTestProjectDoContext(instance));
136
 
137
    context.db.transaction((tx) => {
138
      const projectConfigRow = getProjectConfigRow(tx, input.projectId);
139
      if (!projectConfigRow) {
140
        throw new Error(`Project config ${input.projectId} is missing during test queue fill.`);
141
      }
142
 
143
      // Queue saturation tests only need ProjectDO rows. Skipping eager RunDO
144
      // initialization removes a large amount of setup overhead on busy CI.
145
      const resolvedInput = {
146
        projectId: input.projectId,
147
        triggeredByUserId: input.triggeredByUserId,
148
        branch: input.branch ?? expectTrusted(BranchName, projectConfigRow.defaultBranch, "BranchName"),
149
        repoUrl: projectConfigRow.repoUrl,
150
        configPath: projectConfigRow.configPath,
151
        dispatchMode: expectTrusted(DispatchMode, projectConfigRow.dispatchMode, "DispatchMode"),
152
        executionRuntime: expectTrusted(ExecutionRuntime, projectConfigRow.executionRuntime, "ExecutionRuntime"),
153
      };
154
      const now = Date.now();
155
 
156
      for (let index = 0; index < count; index += 1) {
157
        const accepted = transitionAcceptManualRun(context, tx, resolvedInput, now + index);
158
        if (accepted.kind !== "accepted") {
159
          throw new Error(`Expected accepted manual run during test queue fill, got ${accepted.reason}.`);
160
        }
161
      }
162
    });
163
  });
164
};
165
 
166
export const claimRunWorkWithoutAlarm = async (
167
  projectStub: ReturnType<typeof env.PROJECT_DO.getByName>,
168
  input: Parameters<typeof claimRunWorkCommand>[1],
169
): Promise<Awaited<ReturnType<typeof claimRunWorkCommand>>> => {
170
  let claim: Awaited<ReturnType<typeof claimRunWorkCommand>> | null = null;
171
 
172
  await runInDurableObject(projectStub, async (instance: ProjectDO) => {
173
    const context = createAlarmSchedulingSuppressedContext(createTestProjectDoContext(instance));
174
    claim = await claimRunWorkCommand(context, input);
175
  });
176
 
177
  if (!claim) {
178
    throw new Error("claimRunWorkWithoutAlarm did not return a claim result.");
179
  }
180
 
181
  return claim;
182
};
183
 
184
export const finalizeRunExecutionWithoutAlarm = async (
185
  projectStub: ReturnType<typeof env.PROJECT_DO.getByName>,
186
  input: Omit<Parameters<typeof finalizeRunExecutionCommand>[1], "sandboxDestroyed"> & {
187
    sandboxDestroyed?: boolean;
188
  },
189
): Promise<Awaited<ReturnType<typeof finalizeRunExecutionCommand>>> => {
190
  let result: Awaited<ReturnType<typeof finalizeRunExecutionCommand>> | null = null;
191
 
192
  await runInDurableObject(projectStub, async (instance: ProjectDO) => {
193
    const context = createAlarmSchedulingSuppressedContext(createTestProjectDoContext(instance));
194
    result = await finalizeRunExecutionCommand(context, {
195
      ...input,
196
      sandboxDestroyed: input.sandboxDestroyed ?? true,
197
    });
198
  });
199
 
200
  if (!result) {
201
    throw new Error("finalizeRunExecutionWithoutAlarm did not return a finalize result.");
202
  }
203
 
204
  return result;
205
};
206
 
207
export const requestRunCancelWithoutAlarm = async (
208
  projectStub: ReturnType<typeof env.PROJECT_DO.getByName>,
209
  input: Parameters<typeof requestRunCancelCommand>[1],
210
): Promise<Awaited<ReturnType<typeof requestRunCancelCommand>>> => {
211
  let result: Awaited<ReturnType<typeof requestRunCancelCommand>> | null = null;
212
 
213
  await runInDurableObject(projectStub, async (instance: ProjectDO) => {
214
    const context = createAlarmSchedulingSuppressedContext(createTestProjectDoContext(instance));
215
    result = await requestRunCancelCommand(context, input);
216
  });
217
 
218
  if (!result) {
219
    throw new Error("requestRunCancelWithoutAlarm did not return a cancel result.");
220
  }
221
 
222
  return result;
223
};