1import { describe, expect, it } from "@effect/vitest";
2import * as Effect from "effect/Effect";
3import * as Layer from "effect/Layer";
4import * as Option from "effect/Option";
5import { PgDialect } from "drizzle-orm/pg-core";
6
7import * as RelayDb from "../db.ts";
8import { relayManagedEndpointAllocations } from "../persistence/schema.ts";
9import * as ManagedEndpointAllocations from "./ManagedEndpointAllocations.ts";
10
11const layerWithDb = (db: RelayDb.RelayDb["Service"]) =>
12 ManagedEndpointAllocations.layer.pipe(Layer.provide(Layer.succeed(RelayDb.RelayDb, db)));
13
14describe("ManagedEndpointAllocations", () => {
15 it.effect("clears endpoint readiness and recovery only when the recorded tunnel changes", () => {
16 let updated:
17 | {
18 readonly tunnelId: string;
19 readonly readyAt: unknown;
20 readonly recoveryEnabledAt: unknown;
21 readonly recoveryEnvironmentPublicKey: unknown;
22 readonly tunnelReleasedAt: unknown;
23 }
24 | undefined;
25 const fakeDb = {
26 update: (table: unknown) => {
27 expect(table).toBe(relayManagedEndpointAllocations);
28 return {
29 set: (values: {
30 readonly tunnelId: string;
31 readonly readyAt: unknown;
32 readonly recoveryEnabledAt: unknown;
33 readonly recoveryEnvironmentPublicKey: unknown;
34 readonly tunnelReleasedAt: unknown;
35 }) => {
36 updated = values;
37 return {
38 where: () => ({
39 returning: () => Effect.succeed([{ generation: 8 }]),
40 }),
41 };
42 },
43 };
44 },
45 } as unknown as RelayDb.RelayDb["Service"];
46
47 return Effect.gen(function* () {
48 const allocations = yield* ManagedEndpointAllocations.ManagedEndpointAllocations;
49 expect(
50 yield* allocations.recordTunnel({
51 userId: "user-1",
52 environmentId: "environment-1",
53 tunnelId: "replacement-tunnel",
54 generation: 7,
55 }),
56 ).toBe(8);
57 expect(updated?.tunnelId).toBe("replacement-tunnel");
58 // Any recorded tunnel ends a released state, so status stops reporting it.
59 expect(updated?.tunnelReleasedAt).toBeNull();
60 const query = new PgDialect().sqlToQuery(updated?.readyAt as never);
61 expect(query.sql).toBe(
62 'case when "relay_managed_endpoint_allocations"."tunnel_id" = $1 then "relay_managed_endpoint_allocations"."ready_at" else null end',
63 );
64 expect(query.params).toEqual(["replacement-tunnel"]);
65 for (const [column, value] of [
66 ["recovery_enabled_at", updated?.recoveryEnabledAt],
67 ["recovery_environment_public_key", updated?.recoveryEnvironmentPublicKey],
68 ] as const) {
69 const recoveryQuery = new PgDialect().sqlToQuery(value as never);
70 expect(recoveryQuery.sql).toBe(
71 `case when "relay_managed_endpoint_allocations"."tunnel_id" = $1 then "relay_managed_endpoint_allocations"."${column}" else null end`,
72 );
73 expect(recoveryQuery.params).toEqual(["replacement-tunnel"]);
74 }
75 }).pipe(Effect.provide(layerWithDb(fakeDb)));
76 });
77
78 it.effect("records recovery support and advances the allocation generation", () => {
79 let updated:
80 | {
81 readonly recoveryEnabledAt: string;
82 readonly recoveryEnvironmentPublicKey: string;
83 readonly updatedAt: string;
84 }
85 | undefined;
86 let condition: unknown;
87 const fakeDb = {
88 update: (table: unknown) => {
89 expect(table).toBe(relayManagedEndpointAllocations);
90 return {
91 set: (values: {
92 readonly recoveryEnabledAt: string;
93 readonly recoveryEnvironmentPublicKey: string;
94 readonly updatedAt: string;
95 }) => {
96 updated = values;
97 return {
98 where: (where: unknown) => {
99 condition = where;
100 return {
101 returning: () => Effect.succeed([{ environmentId: "environment-1" }]),
102 };
103 },
104 };
105 },
106 };
107 },
108 } as unknown as RelayDb.RelayDb["Service"];
109
110 return Effect.gen(function* () {
111 const allocations = yield* ManagedEndpointAllocations.ManagedEndpointAllocations;
112 expect(
113 yield* allocations.enableRecovery({
114 userId: "user-1",
115 environmentId: "environment-1",
116 tunnelId: "tunnel-1",
117 environmentPublicKey: "public-key",
118 origin: { localHttpHost: "127.0.0.1", localHttpPort: 3773 },
119 }),
120 ).toBe(true);
121
122 expect(updated?.recoveryEnabledAt).toBe(updated?.updatedAt);
123 expect(updated?.recoveryEnabledAt).toBeDefined();
124 expect(updated?.recoveryEnvironmentPublicKey).toBe("public-key");
125 const query = new PgDialect().sqlToQuery(condition as never);
126 expect(query.sql).toContain('"relay_managed_endpoint_allocations"."tunnel_id"');
127 expect(query.sql).toContain('"relay_environment_links"."environment_public_key"');
128 expect(query.sql).toContain('"relay_environment_links"."endpoint_provider_kind"');
129 expect(query.sql).toContain('"relay_environment_links"."revoked_at" is null');
130 expect(query.sql).toContain("for update");
131 expect(query.params).toContain("tunnel-1");
132 expect(query.params).toContain("public-key");
133 expect(query.params).toContain("cloudflare_tunnel");
134 }).pipe(Effect.provide(layerWithDb(fakeDb)));
135 });
136
137 it.effect("rejects recovery when the tunnel or active link no longer matches", () => {
138 const fakeDb = {
139 update: () => ({
140 set: () => ({
141 where: () => ({
142 returning: () => Effect.succeed([]),
143 }),
144 }),
145 }),
146 } as unknown as RelayDb.RelayDb["Service"];
147
148 return Effect.gen(function* () {
149 const allocations = yield* ManagedEndpointAllocations.ManagedEndpointAllocations;
150 expect(
151 yield* allocations.enableRecovery({
152 userId: "user-1",
153 environmentId: "environment-1",
154 tunnelId: "missing-tunnel",
155 environmentPublicKey: "public-key",
156 origin: { localHttpHost: "127.0.0.1", localHttpPort: 3773 },
157 }),
158 ).toBe(false);
159 }).pipe(Effect.provide(layerWithDb(fakeDb)));
160 });
161
162 it.effect("returns recovery support with tunnel allocation lookups", () => {
163 const base = {
164 userId: "user-1",
165 hostname: "environment.example.test",
166 tunnelName: "managed-tunnel",
167 dnsRecordId: "dns-1",
168 readyAt: "2026-08-25T12:00:00.000Z",
169 updatedAt: "2026-08-25T12:00:00.000Z",
170 generation: 1,
171 };
172 const fakeDb = {
173 select: () => ({
174 from: (table: unknown) => {
175 expect(table).toBe(relayManagedEndpointAllocations);
176 return {
177 leftJoin: () => ({
178 where: () =>
179 Effect.succeed([
180 {
181 ...base,
182 environmentId: "environment-1",
183 tunnelId: "tunnel-1",
184 recoveryEnabledAt: "2026-08-25T12:00:00.000Z",
185 recoveryEnvironmentPublicKey: "current-key",
186 linkedEnvironmentPublicKey: "current-key",
187 },
188 {
189 ...base,
190 environmentId: "environment-2",
191 tunnelId: "tunnel-2",
192 recoveryEnabledAt: null,
193 recoveryEnvironmentPublicKey: null,
194 linkedEnvironmentPublicKey: "current-key",
195 },
196 {
197 ...base,
198 environmentId: "environment-3",
199 tunnelId: "tunnel-3",
200 recoveryEnabledAt: "2026-08-25T12:00:00.000Z",
201 recoveryEnvironmentPublicKey: "old-key",
202 linkedEnvironmentPublicKey: "new-key",
203 },
204 ]),
205 }),
206 };
207 },
208 }),
209 } as unknown as RelayDb.RelayDb["Service"];
210
211 return Effect.gen(function* () {
212 const allocations = yield* ManagedEndpointAllocations.ManagedEndpointAllocations;
213 const result = yield* allocations.listByTunnelNames([
214 "first-tunnel",
215 "second-tunnel",
216 "third-tunnel",
217 ]);
218
219 expect(result.map((entry) => [entry.tunnelId, entry.recoveryEnabled])).toEqual([
220 ["tunnel-1", true],
221 ["tunnel-2", false],
222 ["tunnel-3", false],
223 ]);
224 }).pipe(Effect.provide(layerWithDb(fakeDb)));
225 });
226
227 it.effect("skips the database for an empty tunnel lookup", () =>
228 Effect.gen(function* () {
229 const allocations = yield* ManagedEndpointAllocations.ManagedEndpointAllocations;
230 expect(yield* allocations.listByTunnelNames([])).toEqual([]);
231 }).pipe(Effect.provide(layerWithDb({} as RelayDb.RelayDb["Service"]))),
232 );
233
234 it.effect("splits large tunnel lookups into bounded database queries", () => {
235 const batchSizes: number[] = [];
236 const fakeDb = {
237 select: () => ({
238 from: () => ({
239 leftJoin: () => ({
240 where: (condition: unknown) => {
241 batchSizes.push(new PgDialect().sqlToQuery(condition as never).params.length);
242 return Effect.succeed([]);
243 },
244 }),
245 }),
246 }),
247 } as unknown as RelayDb.RelayDb["Service"];
248
249 return Effect.gen(function* () {
250 const allocations = yield* ManagedEndpointAllocations.ManagedEndpointAllocations;
251 const names = Array.from({ length: 1_001 }, (_, index) => `tunnel-${index}`);
252 expect(yield* allocations.listByTunnelNames(names)).toEqual([]);
253 expect(batchSizes).toEqual([500, 500, 1]);
254 }).pipe(Effect.provide(layerWithDb(fakeDb)));
255 });
256
257 it.effect.each([
258 { markReleased: true, released: true },
259 { markReleased: undefined, released: false },
260 ] as const)(
261 "marks the tunnel released only for the claim that deletes it ($markReleased)",
262 ({ markReleased, released }) => {
263 let updated: Record<string, unknown> | undefined;
264 const fakeDb = {
265 update: () => ({
266 set: (values: Record<string, unknown>) => {
267 updated = values;
268 return {
269 where: () => ({ returning: () => Effect.succeed([{ generation: 4 }]) }),
270 };
271 },
272 }),
273 } as unknown as RelayDb.RelayDb["Service"];
274
275 return Effect.gen(function* () {
276 const allocations = yield* ManagedEndpointAllocations.ManagedEndpointAllocations;
277 yield* allocations.claimRelease({
278 userId: "user-1",
279 environmentId: "environment-1",
280 tunnelId: "tunnel-1",
281 generation: 3,
282 ...(markReleased === undefined ? {} : { markReleased }),
283 });
284 expect(typeof updated?.tunnelReleasedAt === "string").toBe(released);
285 }).pipe(Effect.provide(layerWithDb(fakeDb)));
286 },
287 );
288
289 it.effect("returns a claim generation only when deprovision wins the allocation CAS", () => {
290 const fakeDb = {
291 update: (table: unknown) => {
292 expect(table).toBe(relayManagedEndpointAllocations);
293 return {
294 set: (_values: { readonly updatedAt: string }) => {
295 return {
296 where: () => ({
297 returning: () => Effect.succeed([{ generation: 8 }]),
298 }),
299 };
300 },
301 };
302 },
303 } as unknown as RelayDb.RelayDb["Service"];
304
305 return Effect.gen(function* () {
306 const allocations = yield* ManagedEndpointAllocations.ManagedEndpointAllocations;
307 const generation = yield* allocations.claimDeprovision({
308 userId: "user-1",
309 environmentId: "environment-1",
310 generation: 7,
311 });
312
313 expect(generation).toBe(8);
314 expect(generation).not.toBeNull();
315 }).pipe(Effect.provide(layerWithDb(fakeDb)));
316 });
317
318 it.effect("holds the claimed allocation row while deleting its tunnel", () => {
319 const operations: string[] = [];
320 const fakeDb = {
321 $client: {
322 withTransaction: <A, E, R>(effect: Effect.Effect<A, E, R>) =>
323 Effect.sync(() => {
324 operations.push("transaction");
325 }).pipe(Effect.andThen(effect)),
326 },
327 select: () => ({
328 from: () => ({
329 where: () => ({
330 limit: () => ({
331 for: (strength: string) =>
332 Effect.sync(() => {
333 operations.push(`lock:${strength}`);
334 return [{ generation: 7 }];
335 }),
336 }),
337 }),
338 }),
339 }),
340 } as unknown as RelayDb.RelayDb["Service"];
341
342 return Effect.gen(function* () {
343 const allocations = yield* ManagedEndpointAllocations.ManagedEndpointAllocations;
344 const result = yield* allocations.withClaimedTunnel(
345 {
346 userId: "user-1",
347 environmentId: "environment-1",
348 tunnelId: "tunnel-1",
349 generation: 7,
350 },
351 Effect.sync(() => {
352 operations.push("delete");
353 return true;
354 }),
355 );
356
357 expect(Option.getOrNull(result)).toBe(true);
358 expect(operations).toEqual(["transaction", "lock:update", "delete"]);
359 }).pipe(Effect.provide(layerWithDb(fakeDb)));
360 });
361
362 it.effect("does not remove an allocation superseded after a deprovision claim", () => {
363 const fakeDb = {
364 delete: (table: unknown) => {
365 expect(table).toBe(relayManagedEndpointAllocations);
366 return {
367 where: () => ({
368 returning: () => Effect.succeed([]),
369 }),
370 };
371 },
372 } as unknown as RelayDb.RelayDb["Service"];
373
374 return Effect.gen(function* () {
375 const allocations = yield* ManagedEndpointAllocations.ManagedEndpointAllocations;
376 expect(
377 yield* allocations.removeClaimed({
378 userId: "user-1",
379 environmentId: "environment-1",
380 generation: 7,
381 }),
382 ).toBe(false);
383 }).pipe(Effect.provide(layerWithDb(fakeDb)));
384 });
385
386 it.effect("retains database failures with allocation operation and identity", () => {
387 const cause = new Error("database unavailable");
388 const fakeDb = {
389 select: () => ({
390 from: (table: unknown) => {
391 expect(table).toBe(relayManagedEndpointAllocations);
392 return {
393 where: () => ({
394 limit: () => Effect.fail(cause),
395 }),
396 };
397 },
398 }),
399 } as unknown as RelayDb.RelayDb["Service"];
400
401 return Effect.gen(function* () {
402 const allocations = yield* ManagedEndpointAllocations.ManagedEndpointAllocations;
403 const error = yield* Effect.flip(
404 allocations.get({ userId: "user-1", environmentId: "environment-1" }),
405 );
406
407 expect(error).toMatchObject({
408 _tag: "ManagedEndpointAllocationPersistenceError",
409 operation: "get",
410 stage: "database-request",
411 userId: "user-1",
412 environmentId: "environment-1",
413 });
414 expect(error.cause).toBe(cause);
415 }).pipe(Effect.provide(layerWithDb(fakeDb)));
416 });
417
418 it.effect("reports an unresolved reservation without manufacturing a cause", () => {
419 const fakeDb = {
420 insert: (table: unknown) => {
421 expect(table).toBe(relayManagedEndpointAllocations);
422 return {
423 values: () => ({
424 onConflictDoNothing: () => ({
425 returning: () => Effect.succeed([]),
426 }),
427 }),
428 };
429 },
430 select: () => ({
431 from: (table: unknown) => {
432 expect(table).toBe(relayManagedEndpointAllocations);
433 return {
434 where: () => ({
435 limit: () => Effect.succeed([]),
436 }),
437 };
438 },
439 }),
440 } as unknown as RelayDb.RelayDb["Service"];
441
442 return Effect.gen(function* () {
443 const allocations = yield* ManagedEndpointAllocations.ManagedEndpointAllocations;
444 const error = yield* Effect.flip(
445 allocations.reserve({
446 userId: "user-1",
447 environmentId: "environment-1",
448 hostname: "environment-1.example.test",
449 tunnelName: "environment-1-tunnel",
450 }),
451 );
452
453 expect(error).toMatchObject({
454 _tag: "ManagedEndpointAllocationPersistenceError",
455 operation: "reserve",
456 stage: "resolve-reservation",
457 userId: "user-1",
458 environmentId: "environment-1",
459 hostname: "environment-1.example.test",
460 tunnelName: "environment-1-tunnel",
461 });
462 expect(error.cause).toBeUndefined();
463 expect(error.message).toContain("'resolve-reservation'");
464 }).pipe(Effect.provide(layerWithDb(fakeDb)));
465 });
466});