import { beforeEach, describe, expect, mock, test } from "bun:test"; import { readFileSync } from "node:fs"; import { fakeD1 } from "../metering/fake-d1.test-helper"; let d1 = fakeD1(); const stopped: string[] = []; const fakeEnv: Record = { get DB() { return d1.binding; }, KV: { get: async () => null, put: async () => undefined, delete: async () => undefined }, EDITOR_CONTAINER: { getByName: (name: string) => ({ stop: async () => stopped.push(name) && 0 }) }, METERING_PLANS: "", PUBLIC_URL: "https://forge.test", }; mock.module("cloudflare:workers", () => ({ env: fakeEnv, waitUntil: () => undefined, })); const { closeContainerSession, closeExpiredContainerSessions, closeOrphanSession, openContainerSession } = await import("./container-meter"); const { setPlan } = await import("../metering/quota"); const editorTable = readFileSync(new URL("../../../migrations/0010_editor.sql", import.meta.url), "utf8"); function billed(): number { const row = d1.db.query("SELECT SUM(quantity) AS q FROM usage_events WHERE meter = 'compute.editor'").get() as { q: number | null; } | null; return row?.q ?? 0; } function backdate(id: string, ms: number) { d1.db.query("UPDATE editor_container_sessions SET started_at = started_at - ? WHERE id = ?").run(ms, id); } beforeEach(async () => { d1 = fakeD1(); d1.db.exec(editorTable.slice(0, editorTable.indexOf("CREATE INDEX"))); d1.db.exec( "INSERT INTO accounts (id, name, kind) VALUES ('u1', 'alice', 'user'); INSERT INTO repos VALUES ('r1', 'u1', 'web');", ); await setPlan("u1", "pro", "billing"); }); describe("container sessions", () => { test("a disabled account cannot start a container", async () => { d1.db.exec("UPDATE accounts SET disabled_at = 1, disabled_reason = 'payment_dispute' WHERE id = 'u1'"); await expect(openContainerSession("u1", "r1", "main", 60)).rejects.toThrow("payment dispute"); }); test("bill at most what was held, and a retried close never bills twice", async () => { const hold = await openContainerSession("u1", "r1", "main", 60); backdate(hold.id, 600_000); expect(await closeContainerSession(hold, "stopped")).toBe(60); expect(await closeContainerSession(hold, "stopped")).toBe(60); expect(billed()).toBe(60); }); test("a close retried after a failed commit still bills the original stop time", async () => { const hold = await openContainerSession("u1", "r1", "main", 3600); const stoppedAt = Date.now() + 30_000; d1.db .query("UPDATE editor_container_sessions SET stopped_at = ?, stop_reason = 'idle' WHERE id = ?") .run(stoppedAt, hold.id); const seconds = await closeContainerSession(hold, "stopped", stoppedAt + 600_000); expect(Math.round(seconds)).toBe(30); expect(Math.round(billed())).toBe(30); }); test("cron closes rows whose hold expired, billing only the hold", async () => { const hold = await openContainerSession("u1", "r1", "main", 120); backdate(hold.id, 10 * 3600_000); d1.db.query("UPDATE usage_reservations SET expires_at = ? WHERE id = ?").run(Date.now() - 1, hold.id); expect(await closeExpiredContainerSessions()).toBe(1); expect(stopped).toEqual(["u1:r1:main"]); expect(billed()).toBe(120); const open = d1.db.query("SELECT COUNT(*) AS n FROM editor_container_sessions WHERE stopped_at IS NULL").get(); expect(open).toEqual({ n: 0 }); }); test("orphans with a hold bill at most the hold", async () => { const hold = await openContainerSession("u1", "r1", "main", 90); backdate(hold.id, 3600_000); expect(await closeOrphanSession(hold.id, "stopped")).toBe(90); expect(billed()).toBe(90); }); test("orphans without a hold are capped at the session cap", async () => { const hold = await openContainerSession("u1", "r1", "main", 60); d1.db.query("DELETE FROM usage_reservations WHERE id = ?").run(hold.id); backdate(hold.id, 100 * 3600_000); expect(await closeOrphanSession(hold.id, "stopped")).toBe(8 * 3600); }); });