Shaman
108 lines
4.2 KiB
Python
108 lines
4.2 KiB
Python
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
|