OpenAI-compatible adapter for ComfyUI image generation and editing

2026-10-04 21:56:00 +03:00
commit c17ab5036e
29 changed files with 1730 additions and 0 deletions
@@ -0,0 +1,140 @@
+import base64
+import json
+
+import httpx
+import pytest
+import respx
+
+from tests.conftest import COMFY, png_bytes, png_size
+
+GEN = "http://adapter.test/v1/images/generations"
+PROMPT_ID = "p-1"
+
+
+def _history_json(width=64, height=64):
+    return {
+        PROMPT_ID: {
+            "outputs": {"60": {"images": [{"filename": "out_00001_.png", "subfolder": "", "type": "output"}]}},
+            "status": {"status_str": "success", "completed": True},
+        }
+    }
+
+
+def _mock_pipeline(view_png: bytes | None = None, prompt_response: httpx.Response | None = None):
+    prompt_route = respx.post(f"{COMFY}/prompt").mock(
+        return_value=prompt_response or httpx.Response(200, json={"prompt_id": PROMPT_ID})
+    )
+    respx.get(f"{COMFY}/history/{PROMPT_ID}").mock(return_value=httpx.Response(200, json=_history_json()))
+    view_response = httpx.Response(200, content=view_png) if view_png else httpx.Response(500)
+    respx.get(f"{COMFY}/view").mock(return_value=view_response)
+    return prompt_route
+
+
+@respx.mock
+async def test_generation_success(client: httpx.AsyncClient):
+    png = png_bytes(1024, 1024)
+    prompt_route = _mock_pipeline(view_png=png)
+    r = await client.post(GEN, json={"prompt": "a red cube", "model": "z-image-turbo", "size": "1024x1024", "seed": 123})
+    assert r.status_code == 200
+    body = r.json()
+    assert set(body) == {"created", "data"} and set(body["data"][0]) == {"b64_json"}
+    img = base64.b64decode(body["data"][0]["b64_json"])
+    assert png_size(img) == (1024, 1024)
+    # injected values reached the workflow
+    sent = json.loads(prompt_route.calls.last.request.content.decode())["prompt"]
+    assert sent["83:27"]["inputs"]["text"] == "a red cube"
+    assert sent["83:13"]["inputs"]["width"] == 1024 and sent["83:13"]["inputs"]["height"] == 1024
+    assert sent["83:3"]["inputs"]["seed"] == 123
+
+
+@respx.mock
+async def test_generation_size_custom(client: httpx.AsyncClient):
+    prompt_route = _mock_pipeline(view_png=png_bytes(832, 1216))
+    r = await client.post(GEN, json={"prompt": "x", "size": "832x1216"})
+    assert r.status_code == 200
+    sent = json.loads(prompt_route.calls.last.request.content.decode())["prompt"]
+    assert (sent["83:13"]["inputs"]["width"], sent["83:13"]["inputs"]["height"]) == (832, 1216)
+
+
+@respx.mock
+async def test_generation_seed_random_per_request(client: httpx.AsyncClient):
+    prompt_route = _mock_pipeline(view_png=png_bytes())
+    await client.post(GEN, json={"prompt": "x"})
+    seed_a = json.loads(prompt_route.calls.last.request.content.decode())["prompt"]["83:3"]["inputs"]["seed"]
+    await client.post(GEN, json={"prompt": "x"})
+    seed_b = json.loads(prompt_route.calls.last.request.content.decode())["prompt"]["83:3"]["inputs"]["seed"]
+    assert seed_a != seed_b
+
+
+async def _assert_400(r):
+    assert r.status_code == 400
+    err = r.json()["error"]
+    assert err["type"] == "invalid_request_error" and err["message"]
+
+
+async def test_empty_prompt_400(client):
+    await _assert_400(await client.post(GEN, json={"prompt": "   "}))
+
+
+async def test_n_gt_1_400(client):
+    await _assert_400(await client.post(GEN, json={"prompt": "x", "n": 2}))
+
+
+@pytest.mark.parametrize("bad", ["wide", "1024", "1024x", "0x512"])
+async def test_bad_size_400(client, bad):
+    await _assert_400(await client.post(GEN, json={"prompt": "x", "size": bad}))
+
+
+async def test_unknown_model_400(client):
+    await _assert_400(await client.post(GEN, json={"prompt": "x", "model": "dall-e-3"}))
+
+
+async def test_bad_json_400(client):
+    await _assert_400(await client.post(GEN, content=b"{not json", headers={"Content-Type": "application/json"}))
+
+
+@respx.mock
+async def test_comfy_validation_failure_502(client):
+    _mock_pipeline(prompt_response=httpx.Response(400, json={"node_errors": {"83:27": {"errors": [{"message": "bad", "details": "'83:43'"}]}}}))
+    r = await client.post(GEN, json={"prompt": "x"})
+    assert r.status_code == 502
+    assert r.json()["error"]["code"] == "comfyui_validation_failed"
+
+
+@respx.mock
+async def test_comfy_execution_failure_502(client):
+    respx.post(f"{COMFY}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": PROMPT_ID}))
+    respx.get(f"{COMFY}/history/{PROMPT_ID}").mock(
+        return_value=httpx.Response(200, json={PROMPT_ID: {"outputs": {}, "status": {"status_str": "error", "completed": False}}})
+    )
+    r = await client.post(GEN, json={"prompt": "x"})
+    assert r.status_code == 502
+    assert r.json()["error"]["code"] == "comfyui_execution_failed"
+
+
+@respx.mock
+async def test_comfy_timeout_502(client):
+    respx.post(f"{COMFY}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": PROMPT_ID}))
+    respx.get(f"{COMFY}/history/{PROMPT_ID}").mock(return_value=httpx.Response(200, json={}))
+    # client fixture uses execution_timeout=5; shorten via a second app instance instead:
+    import contextlib
+
+    from app.main import create_app
+    from tests.conftest import settings as make_settings
+
+    app = create_app(make_settings(execution_timeout=0.05, poll_interval=0.0))
+    async with contextlib.AsyncExitStack() as stack:
+        await stack.enter_async_context(app.router.lifespan_context(app))
+        transport = httpx.ASGITransport(app=app)
+        ac = await stack.enter_async_context(httpx.AsyncClient(transport=transport, base_url="http://adapter.test"))
+        ac.headers["Authorization"] = client.headers["Authorization"]
+        r = await ac.post(GEN, json={"prompt": "x"})
+    assert r.status_code == 502
+    assert r.json()["error"]["code"] == "comfyui_timeout"
+
+
+@respx.mock
+async def test_comfy_unreachable_502(client):
+    respx.post(f"{COMFY}/prompt").mock(side_effect=httpx.ConnectError("boom"))
+    r = await client.post(GEN, json={"prompt": "x"})
+    assert r.status_code == 502