OpenAI-compatible adapter for ComfyUI image generation and editing
This commit is contained in:
Shaman
@@ -0,0 +1,51 @@
|
||||
import struct
|
||||
import zlib
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from app.config import Settings
|
||||
from app.main import create_app
|
||||
|
||||
COMFY = "http://comfy.test"
|
||||
FAKE_KEY = "sk-test-1234"
|
||||
|
||||
|
||||
def settings(**overrides) -> Settings:
|
||||
base = dict(
|
||||
comfy_base_url=COMFY,
|
||||
adapter_api_key=FAKE_KEY,
|
||||
validate_templates_on_server=False,
|
||||
poll_interval=0.0,
|
||||
execution_timeout=5.0,
|
||||
)
|
||||
base.update(overrides)
|
||||
return Settings(**base)
|
||||
|
||||
|
||||
def png_bytes(w: int = 64, h: int = 64) -> bytes:
|
||||
def chunk(t: bytes, d: bytes) -> bytes:
|
||||
return struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d))
|
||||
|
||||
ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)
|
||||
iend = b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", ihdr) + chunk(b"IDAT", zlib.compress(b"\x00" * (w * 3 + 1) * h)) + chunk(b"IEND", b"")
|
||||
return iend
|
||||
|
||||
|
||||
def png_size(data: bytes) -> tuple[int, int]:
|
||||
return struct.unpack(">II", data[16:24])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def client():
|
||||
"""Async client bound to the app with lifespan running
|
||||
(templates loaded, comfy client targeting http://comfy.test)."""
|
||||
import contextlib
|
||||
|
||||
app = create_app(settings())
|
||||
async with contextlib.AsyncExitStack() as stack:
|
||||
await stack.enter_async_context(app.router.lifespan_context(app))
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://adapter.test") as ac:
|
||||
ac.headers["Authorization"] = f"Bearer {FAKE_KEY}"
|
||||
yield ac
|
||||
@@ -0,0 +1,26 @@
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from tests.conftest import FAKE_KEY
|
||||
|
||||
|
||||
async def test_missing_key_401(client: httpx.AsyncClient):
|
||||
r = await client.post("/v1/images/generations", json={"prompt": "x"}, headers={"Authorization": ""})
|
||||
assert r.status_code == 401
|
||||
assert r.json()["error"]["code"] == "invalid_api_key"
|
||||
|
||||
|
||||
async def test_wrong_key_401(client: httpx.AsyncClient):
|
||||
r = await client.post("/v1/images/generations", json={"prompt": "x"}, headers={"Authorization": "Bearer sk-wrong"})
|
||||
assert r.status_code == 401
|
||||
|
||||
|
||||
async def test_valid_key_accepted(client: httpx.AsyncClient):
|
||||
r = await client.get("/v1/models")
|
||||
assert r.status_code == 200
|
||||
|
||||
|
||||
async def test_health_open_without_key(client: httpx.AsyncClient):
|
||||
r = await client.get("/health", headers={"Authorization": ""})
|
||||
assert r.status_code == 200
|
||||
assert r.json() == {"status": "ok"}
|
||||
@@ -0,0 +1,120 @@
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
from app.comfy_client import ComfyClient, extract_output_image
|
||||
from app.errors import ComfyUIError
|
||||
|
||||
BASE = "http://comfy.test"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def comfy():
|
||||
client = ComfyClient(BASE)
|
||||
yield client
|
||||
|
||||
|
||||
async def _history_entry(filename="out_001.png", subfolder=""):
|
||||
return {"outputs": {"60": {"images": [{"filename": filename, "subfolder": subfolder, "type": "output"}]}},
|
||||
"status": {"status_str": "success", "completed": True}}
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_submit_prompt_ok(comfy):
|
||||
route = respx.post(f"{BASE}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": "p-1"}))
|
||||
assert await comfy.submit_prompt({"1": {"class_type": "X", "inputs": {}}}) == "p-1"
|
||||
body = route.calls.last.request.content.decode()
|
||||
assert '"client_id"' in body
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_submit_prompt_validation_error(comfy):
|
||||
respx.post(f"{BASE}/prompt").mock(
|
||||
return_value=httpx.Response(400, json={"error": {}, "node_errors": {"3": {"errors": [{"message": "bad link", "details": "'83:43'"}]}}})
|
||||
)
|
||||
with pytest.raises(ComfyUIError, match="83:43"):
|
||||
await comfy.submit_prompt({})
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_submit_prompt_server_error(comfy):
|
||||
respx.post(f"{BASE}/prompt").mock(return_value=httpx.Response(500))
|
||||
with pytest.raises(ComfyUIError, match="HTTP 500"):
|
||||
await comfy.submit_prompt({})
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_upload_image_ok_returns_name(comfy):
|
||||
route = respx.post(f"{BASE}/upload/image").mock(return_value=httpx.Response(200, json={"name": "adapter_x.png", "type": "input"}))
|
||||
name = await comfy.upload_image(b"pngbytes")
|
||||
assert name == "adapter_x.png"
|
||||
assert "multipart" in route.calls.last.request.headers["content-type"]
|
||||
assert b"adapter_" in route.calls.last.request.content
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_upload_image_error(comfy):
|
||||
respx.post(f"{BASE}/upload/image").mock(return_value=httpx.Response(500))
|
||||
with pytest.raises(ComfyUIError):
|
||||
await comfy.upload_image(b"x")
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_wait_for_result_returns_entry(comfy):
|
||||
respx.get(f"{BASE}/history/p-1").mock(return_value=httpx.Response(200, json={"p-1": await _history_entry()}))
|
||||
entry = await comfy.wait_for_result("p-1", timeout=5, interval=0)
|
||||
assert entry["status"]["completed"] is True
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_wait_for_result_execution_failed(comfy):
|
||||
respx.get(f"{BASE}/history/p-1").mock(
|
||||
return_value=httpx.Response(200, json={"p-1": {"outputs": {}, "status": {"status_str": "error", "completed": False}}})
|
||||
)
|
||||
with pytest.raises(ComfyUIError, match="execution failed"):
|
||||
await comfy.wait_for_result("p-1", timeout=5, interval=0)
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_wait_for_result_timeout(comfy):
|
||||
respx.get(f"{BASE}/history/p-1").mock(return_value=httpx.Response(200, json={}))
|
||||
with pytest.raises(ComfyUIError, match="Timed out"):
|
||||
await comfy.wait_for_result("p-1", timeout=0.1, interval=0.01)
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_fetch_view_ok(comfy):
|
||||
respx.get(f"{BASE}/view").mock(return_value=httpx.Response(200, content=b"PNGDATA"))
|
||||
assert await comfy.fetch_view("f.png") == b"PNGDATA"
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_fetch_view_error(comfy):
|
||||
respx.get(f"{BASE}/view").mock(return_value=httpx.Response(404))
|
||||
with pytest.raises(ComfyUIError):
|
||||
await comfy.fetch_view("missing.png")
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_validate_graph_ok_uses_front_false(comfy):
|
||||
route = respx.post(f"{BASE}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": "x"}))
|
||||
await comfy.validate_graph({"1": {"class_type": "X", "inputs": {}}})
|
||||
assert "front=false" in route.calls.last.request.url.query.decode() or "front" in str(route.calls.last.request.url)
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_validate_graph_rejects(comfy):
|
||||
respx.post(f"{BASE}/prompt").mock(return_value=httpx.Response(400, json={"node_errors": {"2": {"errors": [{"message": "missing", "details": ""}]}}}))
|
||||
with pytest.raises(ComfyUIError, match="missing"):
|
||||
await comfy.validate_graph({})
|
||||
|
||||
|
||||
def test_extract_output_image_prefers_output_type():
|
||||
entry = {"outputs": {"77": {"images": [{"filename": "prev.png", "type": "temp"}]},
|
||||
"60": {"images": [{"filename": "real.png", "subfolder": "s", "type": "output"}]}}}
|
||||
assert extract_output_image(entry) == ("real.png", "s")
|
||||
|
||||
|
||||
def test_extract_output_image_missing():
|
||||
with pytest.raises(ComfyUIError, match="no output images"):
|
||||
extract_output_image({"outputs": {}})
|
||||
@@ -0,0 +1,32 @@
|
||||
from app.config import DEFAULT_MAPPING, Settings
|
||||
|
||||
|
||||
def test_defaults():
|
||||
s = Settings(_env_file=None)
|
||||
assert s.comfy_base_url == "http://192.168.0.3:8188"
|
||||
assert s.execution_timeout == 300.0
|
||||
assert s.poll_interval == 1.5
|
||||
assert s.mapping() == DEFAULT_MAPPING
|
||||
|
||||
|
||||
def test_env_override(monkeypatch):
|
||||
monkeypatch.setenv("COMFY_BASE_URL", "http://example:9000")
|
||||
monkeypatch.setenv("ADAPTER_API_KEY", "sk-abc")
|
||||
monkeypatch.setenv("EXECUTION_TIMEOUT", "42")
|
||||
monkeypatch.setenv("POLL_INTERVAL", "0.5")
|
||||
s = Settings(_env_file=None)
|
||||
assert s.comfy_base_url == "http://example:9000"
|
||||
assert s.adapter_api_key == "sk-abc"
|
||||
assert s.execution_timeout == 42.0
|
||||
assert s.poll_interval == 0.5
|
||||
|
||||
|
||||
def test_mapping_file_override(tmp_path):
|
||||
f = tmp_path / "mapping.json"
|
||||
f.write_text('{"create": {"prompt": {"node": "1", "field": "t"}}}', encoding="utf-8")
|
||||
s = Settings(_env_file=None, mapping_file=f)
|
||||
m = s.mapping()
|
||||
assert m["create"]["prompt"] == {"node": "1", "field": "t"}
|
||||
# other sections/keys keep defaults
|
||||
assert m["edit"] == DEFAULT_MAPPING["edit"]
|
||||
assert m["create"]["seed"] == DEFAULT_MAPPING["create"]["seed"]
|
||||
@@ -0,0 +1,107 @@
|
||||
import base64
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
from tests.conftest import COMFY, png_bytes, png_size
|
||||
|
||||
EDITS = "http://adapter.test/v1/images/edits"
|
||||
PROMPT_ID = "p-edit"
|
||||
UPLOADED = "adapter_deadbeef_11223344.png"
|
||||
|
||||
HISTORY = {
|
||||
PROMPT_ID: {
|
||||
"outputs": {"60": {"images": [{"filename": "ed_00001_.png", "subfolder": "", "type": "output"}]}},
|
||||
"status": {"status_str": "success", "completed": True},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _mock_pipeline():
|
||||
upload_route = respx.post(f"{COMFY}/upload/image").mock(return_value=httpx.Response(200, json={"name": UPLOADED, "type": "input"}))
|
||||
prompt_route = 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=HISTORY))
|
||||
respx.get(f"{COMFY}/view").mock(return_value=httpx.Response(200, content=png_bytes(1024, 1024)))
|
||||
return upload_route, prompt_route
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_edits_multipart_success(client: httpx.AsyncClient):
|
||||
upload_route, prompt_route = _mock_pipeline()
|
||||
files = {"image": ("in.png", png_bytes(64, 64), "image/png")}
|
||||
r = await client.post(EDITS, files=files, data={"prompt": "make it night"})
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
img = base64.b64decode(body["data"][0]["b64_json"])
|
||||
assert png_size(img) == (1024, 1024)
|
||||
# uploaded file name landed in LoadImage node; prompt in positive TextEncode
|
||||
upload_body = upload_route.calls.last.request.content
|
||||
assert b"in.png" in upload_body or b"adapter_" in upload_body
|
||||
sent = prompt_route.calls.last.request.content.decode()
|
||||
assert UPLOADED in sent
|
||||
assert "make it night" in sent
|
||||
graph = prompt_route.calls.last.request.content # json bytes
|
||||
assert b'"115:110"' in graph # negative node present with empty prompt
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_edits_json_base64_success(client: httpx.AsyncClient):
|
||||
_, prompt_route = _mock_pipeline()
|
||||
b64 = base64.b64encode(png_bytes(64, 64)).decode()
|
||||
r = await client.post(EDITS, json={"image": b64, "prompt": "add a hat", "model": "qwen-image-edit"})
|
||||
assert r.status_code == 200
|
||||
assert "add a hat" in prompt_route.calls.last.request.content.decode()
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_edits_size_ignored(client: httpx.AsyncClient):
|
||||
_, prompt_route = _mock_pipeline()
|
||||
b64 = base64.b64encode(png_bytes(64, 64)).decode()
|
||||
r = await client.post(EDITS, json={"image": b64, "prompt": "x", "size": "512x512"})
|
||||
assert r.status_code == 200
|
||||
sent = prompt_route.calls.last.request.content.decode()
|
||||
assert "512" not in sent # size never reaches the workflow
|
||||
|
||||
|
||||
async def test_edits_missing_image_400(client):
|
||||
r = await client.post(EDITS, data={"prompt": "x"})
|
||||
assert r.status_code == 400
|
||||
|
||||
|
||||
async def test_edits_empty_image_400(client):
|
||||
files = {"image": ("in.png", b"", "image/png")}
|
||||
r = await client.post(EDITS, files=files, data={"prompt": "x"})
|
||||
assert r.status_code == 400
|
||||
|
||||
|
||||
async def test_edits_garbage_image_400(client):
|
||||
files = {"image": ("in.png", b"not-an-image", "image/png")}
|
||||
r = await client.post(EDITS, files=files, data={"prompt": "x"})
|
||||
assert r.status_code == 400
|
||||
|
||||
|
||||
async def test_edits_bad_base64_400(client):
|
||||
r = await client.post(EDITS, json={"image": "!!!not base64!!!", "prompt": "x"})
|
||||
assert r.status_code == 400
|
||||
|
||||
|
||||
async def test_edits_empty_prompt_400(client):
|
||||
b64 = base64.b64encode(png_bytes(8, 8)).decode()
|
||||
r = await client.post(EDITS, json={"image": b64, "prompt": " "})
|
||||
assert r.status_code == 400
|
||||
|
||||
|
||||
async def test_edits_wrong_model_400(client):
|
||||
b64 = base64.b64encode(png_bytes(8, 8)).decode()
|
||||
r = await client.post(EDITS, json={"image": b64, "prompt": "x", "model": "z-image-turbo"})
|
||||
assert r.status_code == 400
|
||||
|
||||
|
||||
@respx.mock
|
||||
async def test_edits_comfy_failure_502(client):
|
||||
respx.post(f"{COMFY}/upload/image").mock(return_value=httpx.Response(200, json={"name": UPLOADED}))
|
||||
respx.post(f"{COMFY}/prompt").mock(return_value=httpx.Response(400, json={"node_errors": {"x": {"errors": [{"message": "bad graph", "details": ""}]}}}))
|
||||
b64 = base64.b64encode(png_bytes(8, 8)).decode()
|
||||
r = await client.post(EDITS, json={"image": b64, "prompt": "x"})
|
||||
assert r.status_code == 502
|
||||
@@ -0,0 +1,29 @@
|
||||
from app.errors import ComfyUIError, InvalidRequest, Unauthorized, error_body
|
||||
from app.schemas import ImageGenerationRequest
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_error_body_canonical_shape():
|
||||
body = error_body("boom", "server_error", "comfyui_timeout")
|
||||
assert body == {"error": {"message": "boom", "type": "server_error", "param": None, "code": "comfyui_timeout"}}
|
||||
|
||||
|
||||
def test_exception_status_mapping():
|
||||
assert InvalidRequest("x").status == 400
|
||||
assert Unauthorized().status == 401
|
||||
assert ComfyUIError("x").status == 502
|
||||
|
||||
|
||||
def test_generation_request_defaults():
|
||||
r = ImageGenerationRequest(prompt="p")
|
||||
assert r.size == "1024x1024"
|
||||
assert r.n == 1
|
||||
assert r.seed is None
|
||||
assert r.model is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", ["", "abc", "1024", "1024x", "0x10", "10x0"])
|
||||
def test_size_validation(bad):
|
||||
with pytest.raises(ValueError):
|
||||
ImageGenerationRequest(prompt="p", size=bad)
|
||||
@@ -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
|
||||
@@ -0,0 +1,73 @@
|
||||
import pytest
|
||||
|
||||
from app.config import DEFAULT_MAPPING
|
||||
from app.errors import InvalidRequest
|
||||
from app.injector import build_workflow, new_seed, parse_size
|
||||
from app.templates import load_all
|
||||
from app.config import Settings
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def templates():
|
||||
return load_all(Settings(_env_file=None))
|
||||
|
||||
|
||||
def test_parse_size():
|
||||
assert parse_size("1024x1024") == (1024, 1024)
|
||||
assert parse_size("832X1216") == (832, 1216)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad", ["1024", "ax1024", "1024x", "0x100", "1024x-2", "wide"])
|
||||
def test_parse_size_invalid(bad):
|
||||
with pytest.raises(InvalidRequest):
|
||||
parse_size(bad)
|
||||
|
||||
|
||||
def test_seed_random_differs():
|
||||
assert new_seed(None) != new_seed(None)
|
||||
|
||||
|
||||
def test_seed_fixed_and_range():
|
||||
assert new_seed(42) == 42
|
||||
with pytest.raises(InvalidRequest):
|
||||
new_seed(-1)
|
||||
|
||||
|
||||
def test_create_injection(templates):
|
||||
create, _ = templates
|
||||
mapping = DEFAULT_MAPPING["create"]
|
||||
graph = build_workflow(create, mapping, prompt="a cat", seed=77, size="832x1216")
|
||||
assert graph["83:27"]["inputs"]["text"] == "a cat"
|
||||
assert graph["83:13"]["inputs"]["width"] == 832
|
||||
assert graph["83:13"]["inputs"]["height"] == 1216
|
||||
assert graph["83:3"]["inputs"]["seed"] == 77
|
||||
# template untouched
|
||||
assert create["83:27"]["inputs"]["text"] == "\n"
|
||||
assert create["83:13"]["inputs"]["width"] == 1280
|
||||
|
||||
|
||||
def test_edit_injection(templates):
|
||||
_, edit = templates
|
||||
mapping = DEFAULT_MAPPING["edit"]
|
||||
graph = build_workflow(edit, mapping, prompt="make it night", seed=5, image_name="adapter_x.png")
|
||||
assert graph["115:111"]["inputs"]["prompt"] == "make it night"
|
||||
assert graph["115:110"]["inputs"]["prompt"] == ""
|
||||
assert graph["78"]["inputs"]["image"] == "adapter_x.png"
|
||||
assert graph["115:3"]["inputs"]["seed"] == 5
|
||||
|
||||
|
||||
def test_edit_seed_varies_per_call(templates):
|
||||
_, edit = templates
|
||||
mapping = DEFAULT_MAPPING["edit"]
|
||||
seeds = set()
|
||||
for _ in range(3):
|
||||
g = build_workflow(edit, mapping, prompt="p", seed=new_seed(None))
|
||||
seeds.add(g["115:3"]["inputs"]["seed"])
|
||||
assert len(seeds) == 3
|
||||
|
||||
|
||||
def test_unknown_mapping_node_raises(templates):
|
||||
create, _ = templates
|
||||
bad_mapping = {"prompt": {"node": "999", "field": "text"}}
|
||||
with pytest.raises(InvalidRequest, match="999"):
|
||||
build_workflow(create, bad_mapping, prompt="p", seed=1)
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Live tests against the real ComfyUI server (192.168.0.3) via a running adapter.
|
||||
|
||||
Manual run:
|
||||
1. start adapter: uvicorn app.main:app --port 8000
|
||||
2. set env: ADAPTER_LIVE=1 ADAPTER_URL=http://127.0.0.1:8000 ADAPTER_KEY=<ADAPTER_API_KEY>
|
||||
3. pytest -m live
|
||||
"""
|
||||
import base64
|
||||
import os
|
||||
import struct
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
pytestmark = pytest.mark.live
|
||||
|
||||
URL = os.environ.get("ADAPTER_URL", "http://127.0.0.1:8000")
|
||||
KEY = os.environ.get("ADAPTER_KEY", "")
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.live,
|
||||
pytest.mark.skipif(not os.environ.get("ADAPTER_LIVE"), reason="set ADAPTER_LIVE=1 to run live tests"),
|
||||
]
|
||||
|
||||
|
||||
def _client() -> httpx.AsyncClient:
|
||||
return httpx.AsyncClient(base_url=URL, headers={"Authorization": f"Bearer {KEY}"}, timeout=600)
|
||||
|
||||
|
||||
def _png_size(data: bytes) -> tuple[int, int]:
|
||||
return struct.unpack(">II", data[16:24])
|
||||
|
||||
|
||||
def _solid_png(w: int = 128, h: int = 96) -> bytes:
|
||||
import zlib
|
||||
|
||||
def chunk(t, d):
|
||||
return struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d))
|
||||
|
||||
ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)
|
||||
raw = b"".join(b"\x00" + bytes((200, 40, 40)) * w for _ in range(h))
|
||||
return b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", ihdr) + chunk(b"IDAT", zlib.compress(raw)) + chunk(b"IEND", b"")
|
||||
|
||||
|
||||
async def test_generate_1024_with_seed():
|
||||
async with _client() as c:
|
||||
r = await c.post("/v1/images/generations", json={"prompt": "a red cube on white background", "size": "1024x1024", "seed": 20261004})
|
||||
assert r.status_code == 200, r.text
|
||||
img = base64.b64decode(r.json()["data"][0]["b64_json"])
|
||||
assert img.startswith(b"\x89PNG")
|
||||
assert _png_size(img) == (1024, 1024)
|
||||
|
||||
|
||||
async def test_same_seed_same_image():
|
||||
async with _client() as c:
|
||||
payload = {"prompt": "a blue sphere", "size": "1024x1024", "seed": 777}
|
||||
r1 = await c.post("/v1/images/generations", json=payload)
|
||||
r2 = await c.post("/v1/images/generations", json=payload)
|
||||
assert r1.json()["data"][0]["b64_json"] == r2.json()["data"][0]["b64_json"]
|
||||
|
||||
|
||||
async def test_edit_output_is_about_1mp_with_input_aspect():
|
||||
async with _client() as c:
|
||||
files = {"image": ("in.png", _solid_png(128, 96), "image/png")}
|
||||
r = await c.post("/v1/images/edits", files=files, data={"prompt": "make the background deep blue"})
|
||||
assert r.status_code == 200, r.text
|
||||
img = base64.b64decode(r.json()["data"][0]["b64_json"])
|
||||
w, h = _png_size(img)
|
||||
assert 0.9 <= w * h / 1_000_000 <= 1.15 # ~1 MP
|
||||
assert abs(w / h - 128 / 96) < 0.05 # aspect preserved
|
||||
@@ -0,0 +1,11 @@
|
||||
import httpx
|
||||
|
||||
|
||||
async def test_models_list(client: httpx.AsyncClient):
|
||||
r = await client.get("/v1/models")
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert body["object"] == "list"
|
||||
ids = [m["id"] for m in body["data"]]
|
||||
assert ids == ["z-image-turbo", "qwen-image-edit"]
|
||||
assert all(m["object"] == "model" for m in body["data"])
|
||||
@@ -0,0 +1,57 @@
|
||||
import pytest
|
||||
|
||||
from app.config import Settings
|
||||
from app.templates import TemplateError, check_mapping, find_broken_links, load_all
|
||||
|
||||
|
||||
def test_find_broken_links_detects_missing_node():
|
||||
graph = {
|
||||
"1": {"class_type": "X", "inputs": {"m": ["9", 0]}},
|
||||
}
|
||||
errors = find_broken_links(graph)
|
||||
assert len(errors) == 1
|
||||
assert "missing node '9'" in errors[0]
|
||||
assert "node 1" in errors[0]
|
||||
|
||||
|
||||
def test_find_broken_links_ok_on_valid_graph():
|
||||
graph = {
|
||||
"1": {"class_type": "X", "inputs": {"m": ["2", 0]}},
|
||||
"2": {"class_type": "Y", "inputs": {}},
|
||||
}
|
||||
assert find_broken_links(graph) == []
|
||||
|
||||
|
||||
def test_check_mapping_missing_node():
|
||||
graph = {"1": {"class_type": "X", "inputs": {"text": ""}}}
|
||||
mapping = {"create": {"prompt": {"node": "404", "field": "text"}}}
|
||||
errors = check_mapping(graph, mapping["create"], "create")
|
||||
assert any("not found" in e for e in errors)
|
||||
|
||||
|
||||
def test_check_mapping_non_scalar_field():
|
||||
graph = {"1": {"class_type": "X", "inputs": {"text": ["2", 0]}}}
|
||||
mapping = {"create": {"prompt": {"node": "1", "field": "text"}}}
|
||||
errors = check_mapping(graph, mapping["create"], "create")
|
||||
assert any("no scalar input" in e for e in errors)
|
||||
|
||||
|
||||
def test_shipped_templates_load_cleanly():
|
||||
s = Settings(_env_file=None)
|
||||
create, edit = load_all(s)
|
||||
assert create["83:28"]["class_type"] == "UNETLoader"
|
||||
assert edit["115:89"]["class_type"] == "LoraLoaderModelOnly"
|
||||
|
||||
|
||||
def test_broken_template_blocks_startup(tmp_path):
|
||||
bad = tmp_path / "bad.json"
|
||||
bad.write_text('{"1": {"class_type": "X", "inputs": {"m": ["9", 0]}}}', encoding="utf-8")
|
||||
good = tmp_path / "good.json"
|
||||
good.write_text(
|
||||
'{"1": {"class_type": "X", "inputs": {"m": ["2", 0]}}, "2": {"class_type": "Y", "inputs": {}}}',
|
||||
encoding="utf-8",
|
||||
)
|
||||
s = Settings(_env_file=None, create_template_path=bad, edit_template_path=good)
|
||||
with pytest.raises(TemplateError) as exc:
|
||||
load_all(s)
|
||||
assert "missing node '9'" in str(exc.value)
|
||||
Reference in New Issue
Block a user