test: tighten run wait assertions

src/agents/run-wait.test.ts

@@ -15,6 +15,55 @@ import {
   waitForAgentRunAndReadUpdatedAssistantReply,
 } from "./run-wait.js";
 
+type AgentWaitGatewayRequest = {
+  method?: string;
+  params?: {
+    runId?: string;
+    timeoutMs?: unknown;
+  };
+  timeoutMs?: unknown;
+};
+
+function expectNumber(value: unknown, label: string): number {
+  expect(typeof value).toBe("number");
+  if (typeof value !== "number") {
+    throw new Error(`expected ${label} to be a number`);
+  }
+  return value;
+}
+
+function gatewayWaitRequests(): AgentWaitGatewayRequest[] {
+  return callGatewayMock.mock.calls.map(([request]) => request as AgentWaitGatewayRequest);
+}
+
+function requireRequestAt(
+  requests: readonly AgentWaitGatewayRequest[],
+  index: number,
+): AgentWaitGatewayRequest {
+  const request = requests.at(index);
+  if (!request) {
+    throw new Error(`expected gateway request at index ${index}`);
+  }
+  return request;
+}
+
+function expectAgentWaitRequest(
+  request: AgentWaitGatewayRequest,
+  runId: string,
+  maxParamTimeoutMs: number,
+): void {
+  expect(request.method).toBe("agent.wait");
+  expect(request.params?.runId).toBe(runId);
+
+  const paramTimeoutMs = expectNumber(request.params?.timeoutMs, `${runId} param timeoutMs`);
+  const requestTimeoutMs = expectNumber(request.timeoutMs, `${runId} request timeoutMs`);
+  expect(requestTimeoutMs).toBeGreaterThan(0);
+  expect(requestTimeoutMs).toBeLessThanOrEqual(maxParamTimeoutMs + 2_000);
+  expect(paramTimeoutMs).toBe(requestTimeoutMs - 2_000);
+  expect(paramTimeoutMs).toBeGreaterThan(0);
+  expect(paramTimeoutMs).toBeLessThanOrEqual(maxParamTimeoutMs);
+}
+
 describe("readLatestAssistantReply", () => {
   beforeEach(() => {
     callGatewayMock.mockClear();
@@ -287,29 +336,14 @@ describe("waitForAgentRunsToDrain", () => {
       getPendingRunIds: () => activeRunIds,
     });
 
-    expect(result).toEqual({
-      timedOut: false,
-      pendingRunIds: [],
-      deadlineAtMs: expect.any(Number),
-    });
-    expect(callGatewayMock.mock.calls.map((call) => call[0])).toEqual([
-      {
-        method: "agent.wait",
-        params: {
-          runId: "run-1",
-          timeoutMs: expect.any(Number),
-        },
-        timeoutMs: expect.any(Number),
-      },
-      {
-        method: "agent.wait",
-        params: {
-          runId: "run-2",
-          timeoutMs: expect.any(Number),
-        },
-        timeoutMs: expect.any(Number),
-      },
-    ]);
+    expect(result.timedOut).toBe(false);
+    expect(result.pendingRunIds).toStrictEqual([]);
+    expectNumber(result.deadlineAtMs, "deadlineAtMs");
+
+    const requests = gatewayWaitRequests();
+    expect(requests).toHaveLength(2);
+    expectAgentWaitRequest(requireRequestAt(requests, 0), "run-1", 1_000);
+    expectAgentWaitRequest(requireRequestAt(requests, 1), "run-2", 1_000);
   });
 
   it("deduplicates and trims pending run ids", async () => {
@@ -344,15 +378,9 @@ describe("waitForAgentRunsToDrain", () => {
     });
 
     expect(result.timedOut).toBe(false);
-    expect(callGatewayMock.mock.calls.map((call) => call[0])).toEqual([
-      expect.objectContaining({
-        method: "agent.wait",
-        params: expect.objectContaining({ runId: "run-1" }),
-      }),
-      expect.objectContaining({
-        method: "agent.wait",
-        params: expect.objectContaining({ runId: "run-2" }),
-      }),
-    ]);
+    const requests = gatewayWaitRequests();
+    expect(requests).toHaveLength(2);
+    expectAgentWaitRequest(requireRequestAt(requests, 0), "run-1", 1_000);
+    expectAgentWaitRequest(requireRequestAt(requests, 1), "run-2", 1_000);
   });
 });