"""The ComfyUI client, against recorded shapes. Every response body here was read off a real ComfyUI 0.27.0 while this was being written, rather than reconstructed from documentation -- including the one that matters most, `/history` answering with an empty object while the job is still queued. """ from __future__ import annotations import httpx import pytest from lembas.services.images import comfy CONFIG = comfy.Config(base_url="http://comfy.test:8188", timeout=5.0) HISTORY_DONE = { "p1": { "status": {"status_str": "success", "completed": True, "messages": []}, "outputs": { "9": {"images": [{"filename": "LLeMbas_00001_.png", "subfolder": "", "type": "output"}]} }, } } def _handler(routes): """Answer by path, and record what was asked.""" def handle(request: httpx.Request) -> httpx.Response: for path, responder in routes.items(): if request.url.path.startswith(path): return responder(request) return httpx.Response(404, json={"error": "no route"}) return handle # --- Submitting ---------------------------------------------------------------- async def test_a_workflow_is_queued_and_its_id_returned(mock_http): seen: dict = {} def prompt(request): seen["body"] = request.read().decode() return httpx.Response(200, json={"prompt_id": "p1", "number": 3, "node_errors": {}}) mock_http(_handler({"/prompt": prompt})) assert await comfy.submit(CONFIG, {"3": {"class_type": "KSampler"}}) == "p1" assert "KSampler" in seen["body"] assert "client_id" in seen["body"], "ComfyUI keys its progress socket on this" async def test_a_graph_error_names_the_node(mock_http): """ "Invalid prompt" against a twelve-node document says nothing. The commonest cause by far is a checkpoint that does not exist on that machine, and the node id is what points at it.""" mock_http( _handler( { "/prompt": lambda r: httpx.Response( 400, json={ "error": {"type": "prompt_outputs_failed_validation"}, "node_errors": { "4": {"errors": [{"message": "value not in list: ckpt_name"}]} }, }, ) } ) ) with pytest.raises(comfy.ComfyError) as caught: await comfy.submit(CONFIG, {}) assert "node 4" in caught.value.message assert "ckpt_name" in caught.value.message async def test_an_accepted_workflow_with_no_id_is_refused(mock_http): mock_http(_handler({"/prompt": lambda r: httpx.Response(200, json={"number": 1})})) with pytest.raises(comfy.ComfyError): await comfy.submit(CONFIG, {}) async def test_an_unreachable_comfyui_says_so(mock_http): def refuse(request): raise httpx.ConnectError("nope", request=request) mock_http(_handler({"/prompt": refuse})) with pytest.raises(comfy.ComfyError) as caught: await comfy.submit(CONFIG, {}) assert "Could not reach ComfyUI" in caught.value.message # --- Waiting ------------------------------------------------------------------- async def test_an_empty_history_means_not_yet(mock_http, monkeypatch): """`/history/{id}` is empty while the job is queued and gains the whole record when it ends -- so empty is "not yet", not "nothing". Reading it the other way makes every generation fail the instant it is submitted.""" monkeypatch.setattr(comfy, "POLL_INTERVAL", 0.01) calls = {"n": 0} def history(request): calls["n"] += 1 return httpx.Response(200, json=HISTORY_DONE if calls["n"] >= 3 else {}) mock_http(_handler({"/history": history})) refs = await comfy.await_images(CONFIG, "p1") assert calls["n"] == 3, "it kept asking" assert [ref.filename for ref in refs] == ["LLeMbas_00001_.png"] async def test_waiting_gives_up_eventually(mock_http, monkeypatch): monkeypatch.setattr(comfy, "POLL_INTERVAL", 0.01) mock_http(_handler({"/history": lambda r: httpx.Response(200, json={})})) with pytest.raises(comfy.ComfyError) as caught: await comfy.await_images(comfy.Config(base_url=CONFIG.base_url, timeout=0.05), "p1") assert "did not finish" in caught.value.message async def test_a_failed_workflow_is_not_an_empty_one(mock_http, monkeypatch): monkeypatch.setattr(comfy, "POLL_INTERVAL", 0.01) mock_http( _handler( { "/history": lambda r: httpx.Response( 200, json={ "p1": {"status": {"status_str": "error", "completed": True}, "outputs": {}} }, ) } ) ) with pytest.raises(comfy.ComfyError): await comfy.await_images(CONFIG, "p1") async def test_images_are_read_from_every_node(mock_http, monkeypatch): """A template is somebody else's document and may save from a node called anything, or from two of them. Looking for `SaveImage` by name works on the shipped workflow and on nothing else.""" monkeypatch.setattr(comfy, "POLL_INTERVAL", 0.01) mock_http( _handler( { "/history": lambda r: httpx.Response( 200, json={ "p1": { "status": {"status_str": "success", "completed": True}, "outputs": { "12": { "images": [ {"filename": "a.png", "subfolder": "s", "type": "output"} ] }, "40": {"images": [{"filename": "b.png"}]}, }, } }, ) } ) ) refs = await comfy.await_images(CONFIG, "p1") assert [r.filename for r in refs] == ["a.png", "b.png"] assert refs[0].subfolder == "s" # --- Fetching ------------------------------------------------------------------ async def test_the_image_comes_back_as_bytes(mock_http): seen: dict = {} def view(request): seen["params"] = dict(request.url.params) return httpx.Response(200, content=b"\x89PNG\r\n\x1a\n" + b"x" * 100) mock_http(_handler({"/view": view})) payload = await comfy.fetch_image(CONFIG, comfy.Ref("a.png", "sub", "output")) assert payload.startswith(b"\x89PNG") assert seen["params"] == {"filename": "a.png", "subfolder": "sub", "type": "output"} async def test_an_oversized_image_is_refused(mock_http, monkeypatch): """The one place an external service hands back raw bytes that are written to disk. Neither `audio.speak` nor `openai_client` has a cap to copy.""" monkeypatch.setattr(comfy, "MAX_IMAGE_BYTES", 64) mock_http(_handler({"/view": lambda r: httpx.Response(200, content=b"x" * 4096)})) with pytest.raises(comfy.ComfyError): await comfy.fetch_image(CONFIG, comfy.Ref("a.png")) async def test_an_empty_file_is_refused(mock_http): mock_http(_handler({"/view": lambda r: httpx.Response(200, content=b"")})) with pytest.raises(comfy.ComfyError): await comfy.fetch_image(CONFIG, comfy.Ref("a.png")) # --- Freeing ------------------------------------------------------------------- async def test_freeing_never_raises(mock_http): """It runs on the way out of a generation that has already produced its image. Failing the whole tool because a memory hint was refused would be turning a tidy-up into an error.""" def boom(request): raise httpx.ConnectError("gone", request=request) mock_http(_handler({"/free": boom})) await comfy.free(CONFIG) # must not raise async def test_freeing_asks_for_both(mock_http): seen: dict = {} def free(request): seen["body"] = request.read().decode() return httpx.Response(200, json={}) mock_http(_handler({"/free": free})) await comfy.free(CONFIG) assert "unload_models" in seen["body"] and "free_memory" in seen["body"] # --- Discovery ----------------------------------------------------------------- async def test_discovery_reads_the_option_lists(mock_http): """The shape is `{node: {input: {required: {field: [[...values], {meta}]}}}}` -- a list whose first element is the options.""" mock_http( _handler( { "/object_info/CheckpointLoaderSimple": lambda r: httpx.Response( 200, json={ "CheckpointLoaderSimple": { "input": { "required": {"ckpt_name": [["a.safetensors", "b.safetensors"]]} } } }, ), "/object_info/KSampler": lambda r: httpx.Response( 200, json={ "KSampler": { "input": { "required": { "sampler_name": [["euler", "dpmpp_2m"]], "scheduler": [["normal", "karras"]], } } } }, ), } ) ) checkpoints, samplers, schedulers = await comfy.discover(CONFIG) assert checkpoints == ["a.safetensors", "b.safetensors"] assert samplers == ["euler", "dpmpp_2m"] assert schedulers == ["normal", "karras"] async def test_a_custom_node_pack_that_changes_the_shape_is_survived(mock_http): """It is somebody else's schema. An empty list is a page that says "none found"; an exception is a page that says nothing at all.""" mock_http( _handler( { "/object_info": lambda r: httpx.Response(200, json={"Whatever": {"input": {}}}), } ) ) assert await comfy.discover(CONFIG) == ([], [], [])