diff --git a/services/edge-broker/test/broker.test.mjs b/services/edge-broker/test/broker.test.mjs index e3290dd3..538ad3ae 100644 --- a/services/edge-broker/test/broker.test.mjs +++ b/services/edge-broker/test/broker.test.mjs @@ -24,6 +24,20 @@ function waitForClose(socket) { }); } +async function waitFor(predicate, { timeoutMs = 1000, intervalMs = 10, description = "condition" } = {}) { + const deadline = Date.now() + timeoutMs; + + while (Date.now() < deadline) { + if (predicate()) { + return; + } + + await new Promise((resolve) => setTimeout(resolve, intervalMs)); + } + + throw new Error(`Timed out waiting for ${description}`); +} + test("broker dispatches commands to connected agents", async () => { const broker = createBrokerServer({ authMode: "stub", sharedSecret: "secret", commandTimeoutMs: 2000 }); const address = await broker.listen(0); @@ -165,6 +179,7 @@ test("broker closes browser shell sessions immediately when no agent is connecte const browser = new WebSocket(`ws://127.0.0.1:${port}/ws/browser-shell?token=session-token`); await new Promise((resolve) => browser.once("open", resolve)); await waitForClose(browser); + await waitFor(() => closedSessions.length === 1, { description: "offline shell session close callback" }); assert.deepEqual(closedSessions, [{ transcript: "", reason: "agent_offline" }]);