Shaman
141 lines
5.6 KiB
Python
141 lines
5.6 KiB
Python
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
|