fix(e2e): verify suspend control over WebSocket (#122888)

test/scripts/gateway-network-client.test.ts

@@ -12,6 +12,7 @@ import {
   runGatewayNetworkClient,
   runGatewaySuspensionPostRestartClient,
   runGatewaySuspensionPreRestartClient,
+  verifyPreparedSuspensionSocket,
 } from "../../scripts/e2e/lib/gateway-network/client.mts";
 import { readGatewayNetworkClientConnectTimeoutMs } from "../../scripts/e2e/lib/gateway-network/limits.mts";
 import { onceFrame } from "../../scripts/e2e/lib/gateway-network/ws-frames.mts";
@@ -498,4 +499,123 @@ describe("gateway network client", () => {
       }),
     ).toThrow("identify gateway-draining");
   });
+  function createPreparedSocketHarness(responses: GatewayFrame[]) {
+    const frames = [...responses];
+    const requests: Array<{ method: string; params: Record<string, unknown> }> = [];
+    let closeCount = 0;
+    const socket = {
+      close: () => {
+        closeCount += 1;
+      },
+      send: (payload: string) => {
+        const frame = JSON.parse(payload) as {
+          method: string;
+          params: Record<string, unknown>;
+        };
+        requests.push({ method: frame.method, params: frame.params });
+      },
+    };
+    return {
+      get closeCount() {
+        return closeCount;
+      },
+      requests,
+      deps: {
+        onceFrame: async (
+          _ws: unknown,
+          predicate: (frame: GatewayFrame) => boolean,
+          _timeoutMs?: number,
+        ) => {
+          const response = frames.shift();
+          expect(response).toBeDefined();
+          const frame = {
+            type: "res",
+            id: `s${requests.length}`,
+            ...response,
+          };
+          expect(predicate(frame)).toBe(true);
+          return frame;
+        },
+        openSocket: async () => socket,
+        protocolVersion: 1,
+      },
+    };
+  }
+
+  it("uses one authenticated socket for the prepared suspension control lifecycle", async () => {
+    const suspending = {
+      ok: false,
+      error: {
+        code: "UNAVAILABLE",
+        retryable: true,
+        details: { reason: "gateway-suspending", phase: "prepared" },
+      },
+    };
+    const harness = createPreparedSocketHarness([
+      { ok: true },
+      { ok: true, payload: { status: "ready" } },
+      suspending,
+      { ok: false, error: { code: "INVALID_REQUEST" } },
+      { ok: true, payload: { status: "ready" } },
+      { ok: true, payload: { status: "running", resumed: true } },
+      { ok: true, payload: { status: "running", resumed: false } },
+      healthResponse(),
+    ]);
+    await verifyPreparedSuspensionSocket(
+      {
+        deadline: Date.now() + 1_000,
+        suspensionId: "lease-1",
+        token: "test-token",
+        url: "ws://127.0.0.1:12345",
+      },
+      harness.deps,
+    );
+    expect(harness.requests).toEqual([
+      {
+        method: "connect",
+        params: {
+          minProtocol: 1,
+          maxProtocol: 1,
+          client: {
+            id: "cli",
+            displayName: "docker-net-e2e",
+            version: "dev",
+            platform: process.platform,
+            mode: "cli",
+          },
+          caps: [],
+          auth: { token: "test-token" },
+          role: "operator",
+          scopes: ["operator.admin"],
+        },
+      },
+      { method: "gateway.suspend.status", params: { suspensionId: "lease-1" } },
+      { method: "health", params: {} },
+      {
+        method: "gateway.suspend.resume",
+        params: { suspensionId: "lease-1-wrong" },
+      },
+      { method: "gateway.suspend.status", params: { suspensionId: "lease-1" } },
+      { method: "gateway.suspend.resume", params: { suspensionId: "lease-1" } },
+      { method: "gateway.suspend.resume", params: { suspensionId: "lease-1" } },
+      { method: "health", params: {} },
+    ]);
+    expect(harness.closeCount).toBe(1);
+    const invalidHarness = createPreparedSocketHarness([
+      { ok: true },
+      { ok: true, payload: { status: "running" } },
+    ]);
+    await expect(
+      verifyPreparedSuspensionSocket(
+        {
+          deadline: Date.now() + 1_000,
+          suspensionId: "lease-1",
+          token: "test-token",
+          url: "ws://127.0.0.1:12345",
+        },
+        invalidHarness.deps,
+      ),
+    ).rejects.toThrow("prepared suspension must remain ready");
+    expect(invalidHarness.closeCount).toBe(1);
+  });
 });