1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
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
|
