diff options
Diffstat (limited to 'tests')
| -rw-r--r-- | tests/conftest.py | 13 | ||||
| -rw-r--r-- | tests/test_generate.py | 38 |
2 files changed, 50 insertions, 1 deletions
diff --git a/tests/conftest.py b/tests/conftest.py index 7370405..59e3349 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -30,6 +30,8 @@ class FakeComfy: self.interrupts: list[Any] = [] self.drop_ws = False self.hang = False + self.queue_has_prompt = True + self._stalled: tuple[web.WebSocketResponse, str] | None = None self.url = "" self._ws: web.WebSocketResponse | None = None self._ws_ready = asyncio.Event() @@ -75,6 +77,7 @@ class FakeComfy: async def _play(self, ws: web.WebSocketResponse, pid: str) -> None: if self.hang: + self._stalled = (ws, pid) return if self.drop_ws: await ws.close() @@ -82,10 +85,18 @@ class FakeComfy: for msg in self.script or self.success_script(pid): await ws.send_str(json.dumps(msg)) + async def release(self) -> None: + """Finish a stalled (`hang`) job: send the scripted messages now.""" + assert self._stalled is not None + ws, pid = self._stalled + self.hang = False + await self._play(ws, pid) + async def queue(self, request: web.Request) -> web.Response: if request.method == "POST": return web.json_response({}) - return web.json_response({"queue_running": [[0, "pid-1"]], "queue_pending": []}) + running = [[0, "pid-1"]] if self.queue_has_prompt else [] + return web.json_response({"queue_running": running, "queue_pending": []}) async def interrupt(self, request: web.Request) -> web.Response: self.interrupts.append(await request.json()) 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" |