tool: extract helpers

1 files changed, 38 insertions, 0 deletions

diff --git a/tests/test_generate.py b/tests/test_generate.py
index 37dda32..8eaf44b 100644
--- a/tests/test_generate.py
+++ b/tests/test_generate.py
@@ -91,3 +91,41 @@ async def test_cancel_interrupts_backend(
     with pytest.raises(asyncio.CancelledError):
         await task
     assert comfy.interrupts == [{"prompt_id": "pid-1"}]
+
+
+@pytest.fixture
+def fast_idle(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Shrink the 1s WebSocket idle timeout so stall detection takes ms, not seconds."""
+    import comfier
+
+    real_wait_for = asyncio.wait_for
+
+    async def wait_for(aw: Any, timeout: float | None = None) -> Any:
+        return await real_wait_for(aw, 0.01 if timeout == 1.0 else timeout)
+
+    monkeypatch.setattr(comfier.asyncio, "wait_for", wait_for)
+
+
+async def test_stalled_backend_lost_prompt(
+    make_tools: Callable[[], Any], comfy: FakeComfy, fast_idle: None
+) -> None:
+    comfy.hang = True
+    comfy.queue_has_prompt = False  # not in /queue, and /history is empty
+    tools, _ = make_tools()
+    with pytest.raises(ConnectionError, match="lost prompt pid-1"):
+        await asyncio.wait_for(_run(tools), timeout=5)
+
+
+async def test_stalled_backend_still_queued_keeps_waiting(
+    make_tools: Callable[[], Any], comfy: FakeComfy, fast_idle: None
+) -> None:
+    comfy.hang = True  # silent, but /queue still lists the prompt
+    tools, _ = make_tools()
+    task = asyncio.create_task(_run(tools))
+    while not comfy.submitted:
+        await asyncio.sleep(0.01)
+    await asyncio.sleep(0.3)  # many idle ticks and liveness checks pass
+    assert not task.done()
+    await comfy.release()
+    result, _ = await asyncio.wait_for(task, timeout=5)
+    assert result["status"] == "success"