"""Tests for the ComfyUI async HTTP client."""
import httpx
import pytest
import respx
from src.comfyui_client import ComfyUIClient, extract_image_filenames
from src.models import GenerationResult, ModelInfo, QueueStatus
BASE_URL = "http://127.0.0.1:8188"
@pytest.fixture()
def client() -> ComfyUIClient:
"""Return a ComfyUIClient pointed at the test base URL."""
return ComfyUIClient(base_url=BASE_URL)
# --- health_check ---
@respx.mock
async def test_health_check_success(client: ComfyUIClient, mock_comfyui_responses: dict) -> None:
respx.get(f"{BASE_URL}/system_stats").mock(
return_value=httpx.Response(200, json=mock_comfyui_responses["system_stats"])
)
result = await client.health_check()
assert result is True
@respx.mock
async def test_health_check_connection_error(client: ComfyUIClient) -> None:
respx.get(f"{BASE_URL}/system_stats").mock(side_effect=httpx.ConnectError("refused"))
result = await client.health_check()
assert result is False
@respx.mock
async def test_health_check_http_error(client: ComfyUIClient) -> None:
respx.get(f"{BASE_URL}/system_stats").mock(
return_value=httpx.Response(500, text="Internal Server Error")
)
result = await client.health_check()
# status_code != 200, so should return False
assert result is False
# --- submit_workflow ---
@respx.mock
async def test_submit_workflow(client: ComfyUIClient, mock_comfyui_responses: dict) -> None:
respx.post(f"{BASE_URL}/prompt").mock(
return_value=httpx.Response(200, json=mock_comfyui_responses["prompt_response"])
)
prompt_id = await client.submit_workflow({"some": "workflow"})
assert prompt_id == "test-prompt-id-123"
@respx.mock
async def test_submit_workflow_sends_correct_payload(client: ComfyUIClient) -> None:
route = respx.post(f"{BASE_URL}/prompt").mock(
return_value=httpx.Response(200, json={"prompt_id": "abc"})
)
workflow = {"3": {"class_type": "KSampler"}}
await client.submit_workflow(workflow)
request = route.calls.last.request
body = request.content.decode()
import json
payload = json.loads(body)
assert payload["prompt"] == workflow
assert payload["client_id"] == "comfyui-mcp"
@respx.mock
async def test_submit_workflow_connect_error(client: ComfyUIClient) -> None:
respx.post(f"{BASE_URL}/prompt").mock(side_effect=httpx.ConnectError("refused"))
with pytest.raises(httpx.ConnectError, match="Cannot connect"):
await client.submit_workflow({"some": "workflow"})
# --- get_queue ---
@respx.mock
async def test_get_queue_empty(client: ComfyUIClient, mock_comfyui_responses: dict) -> None:
respx.get(f"{BASE_URL}/queue").mock(
return_value=httpx.Response(200, json=mock_comfyui_responses["queue_empty"])
)
status = await client.get_queue()
assert isinstance(status, QueueStatus)
assert status.pending == 0
assert status.running == 0
@respx.mock
async def test_get_queue_with_items(client: ComfyUIClient, mock_comfyui_responses: dict) -> None:
respx.get(f"{BASE_URL}/queue").mock(
return_value=httpx.Response(200, json=mock_comfyui_responses["queue_with_items"])
)
status = await client.get_queue()
assert status.pending == 1
assert status.running == 1
@respx.mock
async def test_get_queue_connect_error(client: ComfyUIClient) -> None:
respx.get(f"{BASE_URL}/queue").mock(side_effect=httpx.ConnectError("refused"))
with pytest.raises(httpx.ConnectError, match="Cannot connect"):
await client.get_queue()
# --- get_history ---
@respx.mock
async def test_get_history_found(client: ComfyUIClient, mock_comfyui_responses: dict) -> None:
prompt_id = "test-prompt-id-123"
respx.get(f"{BASE_URL}/history/{prompt_id}").mock(
return_value=httpx.Response(200, json=mock_comfyui_responses["history_complete"])
)
result = await client.get_history(prompt_id)
assert result is not None
assert "outputs" in result
@respx.mock
async def test_get_history_not_found(client: ComfyUIClient) -> None:
prompt_id = "nonexistent-id"
respx.get(f"{BASE_URL}/history/{prompt_id}").mock(
return_value=httpx.Response(200, json={})
)
result = await client.get_history(prompt_id)
assert result is None
@respx.mock
async def test_get_history_connect_error(client: ComfyUIClient) -> None:
respx.get(f"{BASE_URL}/history/some-id").mock(side_effect=httpx.ConnectError("refused"))
with pytest.raises(httpx.ConnectError, match="Cannot connect"):
await client.get_history("some-id")
# --- poll_until_complete ---
@respx.mock
async def test_poll_until_complete_success(client: ComfyUIClient) -> None:
prompt_id = "poll-test-id"
# First call: no output yet. Second call: output ready.
empty_response = httpx.Response(200, json={})
complete_response = httpx.Response(
200,
json={
prompt_id: {
"outputs": {
"9": {
"images": [
{"filename": "result_00001_.png", "subfolder": "", "type": "output"}
]
}
}
}
},
)
respx.get(f"{BASE_URL}/history/{prompt_id}").mock(
side_effect=[empty_response, complete_response]
)
result = await client.poll_until_complete(prompt_id, timeout=5.0, interval=0.01)
assert isinstance(result, GenerationResult)
assert result.status == "completed"
assert result.prompt_id == prompt_id
assert "result_00001_.png" in result.images
assert result.elapsed_seconds is not None
assert result.elapsed_seconds > 0
@respx.mock
async def test_poll_until_complete_error(client: ComfyUIClient) -> None:
prompt_id = "error-test-id"
error_response = httpx.Response(
200,
json={
prompt_id: {
"status": {
"status_str": "error",
"completed": False,
"messages": [["execution_error", {"message": "Node failed"}]],
},
"outputs": {},
}
},
)
respx.get(f"{BASE_URL}/history/{prompt_id}").mock(return_value=error_response)
result = await client.poll_until_complete(prompt_id, timeout=5.0, interval=0.01)
assert isinstance(result, GenerationResult)
assert result.status == "error"
assert result.prompt_id == prompt_id
assert result.elapsed_seconds is not None
@respx.mock
async def test_poll_until_complete_timeout(client: ComfyUIClient) -> None:
prompt_id = "timeout-test-id"
# Always return empty history
respx.get(f"{BASE_URL}/history/{prompt_id}").mock(
return_value=httpx.Response(200, json={})
)
result = await client.poll_until_complete(prompt_id, timeout=0.05, interval=0.01)
assert isinstance(result, GenerationResult)
assert result.status == "timeout"
assert result.prompt_id == prompt_id
assert result.elapsed_seconds is not None
# --- list_checkpoints ---
@respx.mock
async def test_list_checkpoints(client: ComfyUIClient, mock_comfyui_responses: dict) -> None:
respx.get(f"{BASE_URL}/object_info").mock(
return_value=httpx.Response(200, json=mock_comfyui_responses["object_info"])
)
models = await client.list_checkpoints()
assert len(models) == 2
assert all(isinstance(m, ModelInfo) for m in models)
names = [m.name for m in models]
assert "sd_xl_base_1.0.safetensors" in names
assert "v1-5-pruned.safetensors" in names
assert all(m.type == "checkpoint" for m in models)
@respx.mock
async def test_list_checkpoints_no_loader(client: ComfyUIClient) -> None:
# object_info without CheckpointLoaderSimple key
respx.get(f"{BASE_URL}/object_info").mock(
return_value=httpx.Response(200, json={"SomeOtherNode": {}})
)
models = await client.list_checkpoints()
assert models == []
@respx.mock
async def test_list_checkpoints_connect_error(client: ComfyUIClient) -> None:
respx.get(f"{BASE_URL}/object_info").mock(side_effect=httpx.ConnectError("refused"))
with pytest.raises(httpx.ConnectError, match="Cannot connect"):
await client.list_checkpoints()
# --- get_image_data ---
@respx.mock
async def test_get_image_data(client: ComfyUIClient) -> None:
image_bytes = b"\x89PNG\r\n\x1a\nfake-image-data"
respx.get(f"{BASE_URL}/view").mock(
return_value=httpx.Response(200, content=image_bytes)
)
result = await client.get_image_data("test.png", subfolder="", img_type="output")
assert result == image_bytes
@respx.mock
async def test_get_image_data_connect_error(client: ComfyUIClient) -> None:
respx.get(f"{BASE_URL}/view").mock(side_effect=httpx.ConnectError("refused"))
with pytest.raises(httpx.ConnectError, match="Cannot connect"):
await client.get_image_data("test.png")
# --- extract_image_filenames ---
def testextract_image_filenames_single_node() -> None:
outputs = {
"9": {
"images": [
{"filename": "img_001.png", "subfolder": "", "type": "output"},
{"filename": "img_002.png", "subfolder": "", "type": "output"},
]
}
}
result = extract_image_filenames(outputs)
assert result == ["img_001.png", "img_002.png"]
def testextract_image_filenames_multiple_nodes() -> None:
outputs = {
"9": {
"images": [{"filename": "a.png", "subfolder": "", "type": "output"}]
},
"10": {
"images": [{"filename": "b.png", "subfolder": "", "type": "output"}]
},
}
result = extract_image_filenames(outputs)
assert "a.png" in result
assert "b.png" in result
assert len(result) == 2
def testextract_image_filenames_empty() -> None:
result = extract_image_filenames({})
assert result == []
def testextract_image_filenames_missing_filename_key() -> None:
outputs = {"9": {"images": [{"subfolder": "", "type": "output"}]}}
result = extract_image_filenames(outputs)
assert result == []
# --- context manager ---
@respx.mock
async def test_async_context_manager(mock_comfyui_responses: dict) -> None:
respx.get(f"{BASE_URL}/system_stats").mock(
return_value=httpx.Response(200, json=mock_comfyui_responses["system_stats"])
)
async with ComfyUIClient(base_url=BASE_URL) as client:
result = await client.health_check()
assert result is True
