claude: improved dev-setup

4 files changed, 296 insertions, 0 deletions

diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/tests/__init__.py
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
diff --git a/tests/test_generate.py b/tests/test_generate.py
new file mode 100644
index 0000000..37dda32
--- /dev/null
+++ b/tests/test_generate.py
@@ -0,0 +1,93 @@
+import asyncio
+from collections.abc import Callable
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+from tests.conftest import FakeComfy
+
+
+async def _run(tools: Any, **kw: Any) -> tuple[dict[str, Any], list[dict[str, Any]]]:
+    events: list[dict[str, Any]] = []
+
+    async def emit(e: dict[str, Any]) -> None:
+        events.append(e)
+
+    result = await tools.comfier_generate_image(
+        positive_prompt="a marble",
+        seed="5",
+        size="64x64",
+        __request__=SimpleNamespace(),
+        __user__={"id": "u"},
+        __event_emitter__=emit,
+        **kw,
+    )
+    return result, events
+
+
+async def test_success(make_tools: Callable[[], Any], comfy: FakeComfy) -> None:
+    tools, env = make_tools()
+    result, events = await _run(tools)
+    assert result["status"] == "success"
+    assert len(result["images"]) == 1 and len(env.stored) == 1
+    assert comfy.submitted[0]["520:517"]["inputs"]["seed"] == 5
+    assert comfy.submitted[0]["520:518"]["inputs"]["value"] == 64
+    assert events[-1]["type"] == "files"
+
+
+async def test_execution_error(make_tools: Callable[[], Any], comfy: FakeComfy) -> None:
+    comfy.script = [
+        {
+            "type": "execution_error",
+            "data": {"prompt_id": "pid-1", "exception_message": "boom"},
+        }
+    ]
+    tools, _ = make_tools()
+    with pytest.raises(RuntimeError, match="boom"):
+        await _run(tools)
+
+
+async def test_ws_drop_retry_declined(
+    make_tools: Callable[[], Any], comfy: FakeComfy
+) -> None:
+    comfy.drop_ws = True
+    tools, _ = make_tools()
+    asked: list[dict[str, Any]] = []
+
+    async def call(e: dict[str, Any]) -> bool:
+        asked.append(e)
+        return False
+
+    with pytest.raises(ConnectionError):
+        await _run(tools, __event_call__=call)
+    assert len(asked) == 1
+
+
+async def test_ws_drop_retry_accepted(
+    make_tools: Callable[[], Any], comfy: FakeComfy
+) -> None:
+    comfy.drop_ws = True
+    tools, _ = make_tools()
+
+    async def call(e: dict[str, Any]) -> bool:
+        comfy.drop_ws = False
+        return True
+
+    result, _ = await _run(tools, __event_call__=call)
+    assert result["status"] == "success"
+
+
+async def test_cancel_interrupts_backend(
+    make_tools: Callable[[], Any], comfy: FakeComfy
+) -> None:
+    comfy.hang = True
+    tools, _ = make_tools()
+    task = asyncio.create_task(_run(tools))
+    while not comfy.submitted:
+        await asyncio.sleep(0.05)
+    await asyncio.sleep(0.1)
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+    assert comfy.interrupts == [{"prompt_id": "pid-1"}]
diff --git a/tests/test_helpers.py b/tests/test_helpers.py
new file mode 100644
index 0000000..0c0f875
--- /dev/null
+++ b/tests/test_helpers.py
@@ -0,0 +1,59 @@
+import json
+import pathlib
+
+import pytest
+
+from comfier import Tools
+
+ROOT = pathlib.Path(__file__).parent.parent
+
+
+@pytest.fixture
+def tools(make_tools):  # type: ignore[no-untyped-def]
+    return make_tools()[0]
+
+
+def test_normalize_seed(tools: Tools) -> None:
+    assert tools._normalize_seed("42") == 42
+    assert tools._normalize_seed("0x10") == 16
+    assert tools._normalize_seed(7) == 7
+    assert 0 <= tools._normalize_seed(None) < 1 << 63
+    assert 0 <= tools._normalize_seed(-5) < 1 << 63
+    assert tools._normalize_seed((1 << 63) + 3) == 3
+    with pytest.raises(ValueError):
+        tools._normalize_seed("abc")
+
+
+async def test_resolve_dimensions(tools: Tools) -> None:
+    assert await tools._resolve_dimensions("512X768") == (512, 768)
+    assert await tools._resolve_dimensions(None) == (1024, 1024)
+    with pytest.raises(ValueError):
+        await tools._resolve_dimensions("0x10")
+    with pytest.raises(ValueError):
+        await tools._resolve_dimensions("bogus")
+
+
+def test_resolve_node_mappings(tools: Tools) -> None:
+    tools.valves.negative_prompt_node_ids = "1, 2,,"
+    nodes = json.loads((ROOT / "fixtures" / "nodes.json").read_text())
+    m = tools._resolve_node_mappings(nodes)
+    assert m["negative_prompt"]["node_ids"] == ["1", "2"]
+    assert m["positive_prompt"] == {"key": "value", "node_ids": ["520:506"]}
+    assert m["steps"]["node_ids"] == ["459:458"]
+
+
+def test_apply_workflow_parameters(tools: Tools) -> None:
+    wf = json.loads((ROOT / "fixtures" / "workflow.json").read_text())
+    m = tools._resolve_node_mappings(
+        json.loads((ROOT / "fixtures" / "nodes.json").read_text())
+    )
+    applied = tools._apply_workflow_parameters(
+        wf, m, {"seed": 9, "width": 64, "steps": None}
+    )
+    assert wf["520:517"]["inputs"]["seed"] == 9
+    assert wf["520:518"]["inputs"]["value"] == 64
+    assert "steps" not in applied
+    with pytest.raises(KeyError):
+        tools._apply_workflow_parameters(
+            wf, {"seed": {"key": "seed", "node_ids": ["nope"]}}, {"seed": 1}
+        )