cloudflare-container.test.ts

// CloudflareContainerBackend tests — exercise the lifecycle
// plumbing against an in-process fake IWorkspaceContainerAPI.
//
// The successful connect() path constructs a WebSocketPair, which
// is a workerd global not available under the vitest node runner.
// These tests cover the paths that bail before the upgrade (port
// never opens, /connect non-2xx, /ws upgrade timeout), the
// handleFetch input validation, and the factory + workspace-ref
// plumbing. The full happy-path round-trip is covered by the live
// example.

import { describe, expect, test, vi } from "vitest";

import { CloudflareContainerBackend } from "./cloudflare-container.js";
import type { IWorkspaceContainerAPI, WorkspaceRef } from "./container-host.js";

interface FakeHostOptions {
  healthy?: boolean;
  // Health probe sequence: each connect() reads from the head of
  // this array. true = answer 200, false = throw "connection
  // refused". A single `healthy` flag still works for tests that
  // don't care about transitions.
  healthSequence?: boolean[];
  connectStatus?: number;
  restart?: () => Promise<void>;
  // Pre-set a prior exit reason so connect()'s pre-flight
  // exitInfo() check observes it.
  priorExit?: { exitedAt: number; reason: string } | null;
}

interface FakeHost {
  host: IWorkspaceContainerAPI;
  calls: { name: string; args: unknown[] }[];
  startEnv?: Record<string, string>;
  interceptedHost?: string;
  interceptedWorkspace?: WorkspaceRef;
  running: boolean;
  exit: { exitedAt: number; reason: string } | null;
  simulateExit(reason: string): void;
}

function makeFakeHost(opts: FakeHostOptions = {}): FakeHost {
  const healthSequence = opts.healthSequence?.slice();
  const defaultHealthy = opts.healthy ?? true;
  const connectStatus = opts.connectStatus ?? 200;
  const calls: { name: string; args: unknown[] }[] = [];
  const state: FakeHost = {
    calls,
    running: false,
    exit: opts.priorExit ?? null,
    simulateExit(reason: string) {
      state.exit = { exitedAt: Date.now(), reason };
      state.running = false;
    },
  } as FakeHost;

  function nextHealthy(): boolean {
    if (healthSequence && healthSequence.length > 0) {
      return healthSequence.shift() ?? defaultHealthy;
    }
    return defaultHealthy;
  }

  state.host = {
    async start(env) {
      calls.push({ name: "start", args: [env] });
      state.startEnv = env;
      state.running = true;
      // A successful start clears any prior exit, matching
      // WorkspaceContainerAPI.start.
      state.exit = null;
    },
    async interceptOutboundHttp(host, ref) {
      calls.push({ name: "interceptOutboundHttp", args: [host, ref] });
      state.interceptedHost = host;
      state.interceptedWorkspace = ref;
    },
    async fetchPort(port, input, init) {
      const request = input instanceof Request ? input : new Request(input, init);
      const url = new URL(request.url);
      calls.push({ name: "fetchPort", args: [port, url.pathname, request.method] });
      if (url.pathname === "/health") {
        if (!nextHealthy()) throw new Error("connection refused");
        return new Response(null, { status: 200 });
      }
      if (url.pathname === "/connect") {
        if (connectStatus !== 200) {
          return new Response(`/connect ${connectStatus}`, { status: connectStatus });
        }
        return new Response(JSON.stringify({ ok: true }), { status: 200 });
      }
      throw new Error(`unexpected port path: ${url.pathname}`);
    },
    port() {
      throw new Error("cross-boundary Fetchers should not be used by CloudflareContainerBackend");
    },
    async restart(env) {
      calls.push({ name: "restart", args: [env] });
      if (opts.restart) {
        await opts.restart();
      }
      state.running = true;
      state.exit = null;
    },
    async status() {
      calls.push({ name: "status", args: [] });
      return { running: state.running, exit: state.exit };
    },
    async exitInfo() {
      calls.push({ name: "exitInfo", args: [] });
      return state.exit;
    },
  } satisfies IWorkspaceContainerAPI;
  return state;
}

const fakeWorkspace: WorkspaceRef = { binding: "TestDO", id: "abc123" };

describe("CloudflareContainerBackend", () => {
  test("connect() throws when the container port never opens", async () => {
    const fake = makeFakeHost({ healthy: false });
    const backend = new CloudflareContainerBackend({
      container: () => ({ getWorkspaceContainer: () => fake.host }),
      workspace: fakeWorkspace,
      connectTimeoutMs: 600,
      restartAttempts: 0,
    });

    await expect(backend.connect()).rejects.toThrow(/stage=health.*port=8080/);

    const names = fake.calls.map((c) => c.name);
    expect(names).toContain("start");
    expect(names).toContain("interceptOutboundHttp");
    expect(fake.interceptedHost).toBe("workspace.internal");
    expect(fake.interceptedWorkspace).toEqual(fakeWorkspace);
  });

  test("egressHost option overrides the default", async () => {
    const fake = makeFakeHost({ healthy: false });
    const backend = new CloudflareContainerBackend({
      container: () => ({ getWorkspaceContainer: () => fake.host }),
      workspace: fakeWorkspace,
      egressHost: "wsd.local",
      connectTimeoutMs: 300,
    });
    await expect(backend.connect()).rejects.toThrow();
    expect(fake.interceptedHost).toBe("wsd.local");
  });

  test("containerEnv option merges onto the start() env", async () => {
    const fake = makeFakeHost({ healthy: false });
    const backend = new CloudflareContainerBackend({
      container: () => ({ getWorkspaceContainer: () => fake.host }),
      workspace: fakeWorkspace,
      containerEnv: { CUSTOM: "1", PORT: "9000" },
      connectTimeoutMs: 300,
    });
    await expect(backend.connect()).rejects.toThrow();
    expect(fake.startEnv?.CUSTOM).toBe("1");
    // Caller-supplied value wins over the default.
    expect(fake.startEnv?.PORT).toBe("9000");
    // Defaults still flow through.
    expect(fake.startEnv?.MOUNT_POINT).toBe("/workspace");
  });

  test("container factory is invoked per connect()", async () => {
    const fake = makeFakeHost({ healthy: false });
    const factory = vi.fn(() => ({ getWorkspaceContainer: () => fake.host }));
    const backend = new CloudflareContainerBackend({
      container: factory,
      workspace: fakeWorkspace,
      connectTimeoutMs: 300,
    });
    await expect(backend.connect()).rejects.toThrow();
    await expect(backend.connect()).rejects.toThrow();
    // Two failed dials → two factory invocations. The cached
    // handle only short-circuits on success.
    expect(factory).toHaveBeenCalledTimes(2);
  });

  test("async container factory is awaited", async () => {
    const fake = makeFakeHost({ healthy: false });
    const backend = new CloudflareContainerBackend({
      container: async () => {
        await Promise.resolve();
        return { getWorkspaceContainer: () => fake.host };
      },
      workspace: fakeWorkspace,
      connectTimeoutMs: 300,
    });
    await expect(backend.connect()).rejects.toThrow();
    expect(fake.calls.map((c) => c.name)).toContain("start");
  });

  test("connect() throws when /connect returns non-2xx", async () => {
    const fake = makeFakeHost({ connectStatus: 502 });
    const backend = new CloudflareContainerBackend({
      container: () => ({ getWorkspaceContainer: () => fake.host }),
      workspace: fakeWorkspace,
      connectTimeoutMs: 600,
    });

    await expect(backend.connect()).rejects.toThrow(/POST \/connect returned 502/);
  });

  test("connect() throws when the /ws upgrade never arrives", async () => {
    const fake = makeFakeHost();
    const backend = new CloudflareContainerBackend({
      container: () => ({ getWorkspaceContainer: () => fake.host }),
      workspace: fakeWorkspace,
      connectTimeoutMs: 600,
    });

    await expect(backend.connect()).rejects.toThrow(/\/ws upgrade did not arrive/);
  });

  test("handleFetch rejects non-/ws paths", async () => {
    const fake = makeFakeHost();
    const backend = new CloudflareContainerBackend({
      container: () => ({ getWorkspaceContainer: () => fake.host }),
      workspace: fakeWorkspace,
    });
    const res = await backend.handleFetch(new Request("http://workspace.internal/other"));
    expect(res.status).toBe(404);
  });

  test("handleFetch rejects missing upgrade header", async () => {
    const fake = makeFakeHost();
    const backend = new CloudflareContainerBackend({
      container: () => ({ getWorkspaceContainer: () => fake.host }),
      workspace: fakeWorkspace,
    });
    const res = await backend.handleFetch(new Request("http://workspace.internal/ws"));
    expect(res.status).toBe(426);
  });

  test("connect() consults host.exitInfo() before host.start()", async () => {
    const fake = makeFakeHost();
    const backend = new CloudflareContainerBackend({
      container: () => ({ getWorkspaceContainer: () => fake.host }),
      workspace: fakeWorkspace,
      connectTimeoutMs: 600,
      restartAttempts: 0,
    });
    // Doesn't matter that this rejects — we just want to observe
    // the call order.
    await backend.connect().catch(() => undefined);
    const names = fake.calls.map((c) => c.name);
    const exitIdx = names.indexOf("exitInfo");
    const startIdx = names.indexOf("start");
    expect(exitIdx).toBeGreaterThanOrEqual(0);
    expect(startIdx).toBeGreaterThanOrEqual(0);
    expect(exitIdx).toBeLessThan(startIdx);
  });

  test("connect() surfaces a prior exit reason in the stage-tagged error", async () => {
    const fake = makeFakeHost({
      healthy: false,
      priorExit: { exitedAt: Date.now() - 5_000, reason: "OOM killed" },
    });
    const backend = new CloudflareContainerBackend({
      container: () => ({ getWorkspaceContainer: () => fake.host }),
      workspace: fakeWorkspace,
      connectTimeoutMs: 600,
      restartAttempts: 0,
    });
    const err = await backend.connect().then(
      () => undefined,
      (e: Error) => e,
    );
    expect(String(err)).toMatch(/stage=health/);
    expect(String(err)).toMatch(/priorExit="OOM killed"/);
  });

  test("connect() restarts the host when initial readiness fails and recovers", async () => {
    // First attempt drains all probes as failures; restart() runs;
    // the second attempt's very first probe answers healthy.
    // connect() still fails at the /ws upgrade (no WebSocketPair
    // under node) — the point is that readiness recovered after
    // restart and we reached the /connect POST and /ws upgrade.
    const fake = makeFakeHost({
      healthSequence: [
        // First attempt — enough failures to exhaust the budget.
        false,
        false,
        false,
        false,
        false,
        // Restart, then second attempt: first probe is healthy.
        true,
      ],
    });
    const backend = new CloudflareContainerBackend({
      container: () => ({ getWorkspaceContainer: () => fake.host }),
      workspace: fakeWorkspace,
      connectTimeoutMs: 2000,
      restartAttempts: 1,
    });
    await expect(backend.connect()).rejects.toThrow(/stage=ws/);
    const names = fake.calls.map((c) => c.name);
    expect(names.filter((n) => n === "start")).toHaveLength(1);
    expect(names.filter((n) => n === "restart")).toHaveLength(1);
    // /connect was reached after restart succeeded.
    const paths = fake.calls.filter((c) => c.name === "fetchPort").map((c) => c.args[1] as string);
    expect(paths).toContain("/connect");
  });

  test("connect() surfaces stage='health' when readiness exhausts all attempts", async () => {
    const fake = makeFakeHost({ healthy: false });
    const backend = new CloudflareContainerBackend({
      container: () => ({ getWorkspaceContainer: () => fake.host }),
      workspace: fakeWorkspace,
      connectTimeoutMs: 800,
      restartAttempts: 1,
    });
    const err = await backend.connect().then(
      () => undefined,
      (e: Error) => e,
    );
    expect(err).toBeDefined();
    const msg = String(err);
    expect(msg).toMatch(/stage=health/);
    expect(msg).toMatch(/attempts?=2/);
    expect(msg).toMatch(/port=8080/);
    // restart was attempted before giving up.
    expect(fake.calls.some((c) => c.name === "restart")).toBe(true);
  });

  test("connect() reports stage='health' when restartAttempts=0 and probe never succeeds", async () => {
    const fake = makeFakeHost({ healthy: false });
    const backend = new CloudflareContainerBackend({
      container: () => ({ getWorkspaceContainer: () => fake.host }),
      workspace: fakeWorkspace,
      connectTimeoutMs: 600,
      restartAttempts: 0,
    });
    const err = await backend.connect().then(
      () => undefined,
      (e: Error) => e,
    );
    expect(String(err)).toMatch(/stage=health/);
    // No restart attempt.
    expect(fake.calls.some((c) => c.name === "restart")).toBe(false);
  });
});