tool: extract helpers

1 files changed, 12 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())