1 files changed, 144 insertions, 0 deletions
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..7370405
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,144 @@
+import asyncio
+import json
+import pathlib
+from collections.abc import AsyncIterator, Callable
+from typing import Any
+
+import pytest
+import pytest_asyncio
+from aiohttp import WSMsgType, web
+
+from devkit import FakeEnv, install
+
+ROOT = pathlib.Path(__file__).parent.parent
+PNG = bytes.fromhex(
+    "89504e470d0a1a0a0000000d49484452000000010000000108060000001f15c489"
+    "0000000d49444154789c6360000002000001e221bc330000000049454e44ae426082"
+)
+
+
+# Make `import comfier` work at collection time; fixtures re-install per test.
+install(FakeEnv(comfyui_url="", workflow="{}"))
+
+
+class FakeComfy:
+    """Minimal in-process ComfyUI. `script` lists WS messages sent after /prompt."""
+
+    def __init__(self) -> None:
+        self.script: list[dict[str, Any]] = []
+        self.submitted: list[dict[str, Any]] = []
+        self.interrupts: list[Any] = []
+        self.drop_ws = False
+        self.hang = False
+        self.url = ""
+        self._ws: web.WebSocketResponse | None = None
+        self._ws_ready = asyncio.Event()
+
+    def success_script(self, prompt_id: str) -> list[dict[str, Any]]:
+        return [
+            {"type": "executing", "data": {"node": "1", "prompt_id": prompt_id}},
+            {
+                "type": "progress",
+                "data": {"value": 1, "max": 2, "prompt_id": prompt_id},
+            },
+            {
+                "type": "executed",
+                "data": {
+                    "prompt_id": prompt_id,
+                    "output": {"images": [{"filename": "a.png", "type": "output"}]},
+                },
+            },
+            {"type": "executing", "data": {"node": None, "prompt_id": prompt_id}},
+        ]
+
+    async def ws(self, request: web.Request) -> web.WebSocketResponse:
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+        self._ws = ws
+        self._ws_ready.set()
+        async for msg in ws:
+            if msg.type == WSMsgType.ERROR:
+                break
+        return ws
+
+    async def prompt(self, request: web.Request) -> web.Response:
+        body = await request.json()
+        self.submitted.append(body["prompt"])
+        pid = "pid-1"
+        await self._ws_ready.wait()
+        assert self._ws is not None
+        ws = self._ws
+        asyncio.get_running_loop().call_soon(
+            lambda: asyncio.ensure_future(self._play(ws, pid))
+        )
+        return web.json_response({"prompt_id": pid})
+
+    async def _play(self, ws: web.WebSocketResponse, pid: str) -> None:
+        if self.hang:
+            return
+        if self.drop_ws:
+            await ws.close()
+            return
+        for msg in self.script or self.success_script(pid):
+            await ws.send_str(json.dumps(msg))
+
+    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": []})
+
+    async def interrupt(self, request: web.Request) -> web.Response:
+        self.interrupts.append(await request.json())
+        return web.json_response({})
+
+    async def history(self, request: web.Request) -> web.Response:
+        return web.json_response({})
+
+    async def view(self, request: web.Request) -> web.Response:
+        return web.Response(body=PNG, content_type="image/png")
+
+
+@pytest_asyncio.fixture
+async def comfy() -> AsyncIterator[FakeComfy]:
+    fake = FakeComfy()
+    app = web.Application()
+    app.add_routes(
+        [
+            web.get("/ws", fake.ws),
+            web.post("/prompt", fake.prompt),
+            web.get("/queue", fake.queue),
+            web.post("/queue", fake.queue),
+            web.post("/interrupt", fake.interrupt),
+            web.get("/history/{id}", fake.history),
+            web.get("/view", fake.view),
+        ]
+    )
+    runner = web.AppRunner(app)
+    await runner.setup()
+    site = web.TCPSite(runner, "127.0.0.1", 0)
+    await site.start()
+    port = site._server.sockets[0].getsockname()[1]  # type: ignore[union-attr]
+    fake.url = f"http://127.0.0.1:{port}"
+    yield fake
+    await runner.cleanup()
+
+
+@pytest.fixture
+def make_tools(
+    tmp_path: pathlib.Path, comfy: FakeComfy
+) -> Callable[[], tuple[Any, FakeEnv]]:
+    def factory() -> tuple[Any, FakeEnv]:
+        env = FakeEnv(
+            comfyui_url=comfy.url,
+            workflow=(ROOT / "fixtures" / "workflow.json").read_text(encoding="utf-8"),
+            nodes=json.loads(
+                (ROOT / "fixtures" / "nodes.json").read_text(encoding="utf-8")
+            ),
+            output_dir=tmp_path,
+        )
+        install(env)
+        from comfier import Tools
+
+        return Tools(), env
+
+    return factory