Shaman
121 lines
4.4 KiB
Python
121 lines
4.4 KiB
Python
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": {}})
|