1 files changed, 93 insertions, 0 deletions
diff --git a/tests/test_generate.py b/tests/test_generate.py
new file mode 100644
index 0000000..37dda32
--- /dev/null
+++ b/tests/test_generate.py
@@ -0,0 +1,93 @@
+import asyncio
+from collections.abc import Callable
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+from tests.conftest import FakeComfy
+
+
+async def _run(tools: Any, **kw: Any) -> tuple[dict[str, Any], list[dict[str, Any]]]:
+    events: list[dict[str, Any]] = []
+
+    async def emit(e: dict[str, Any]) -> None:
+        events.append(e)
+
+    result = await tools.comfier_generate_image(
+        positive_prompt="a marble",
+        seed="5",
+        size="64x64",
+        __request__=SimpleNamespace(),
+        __user__={"id": "u"},
+        __event_emitter__=emit,
+        **kw,
+    )
+    return result, events
+
+
+async def test_success(make_tools: Callable[[], Any], comfy: FakeComfy) -> None:
+    tools, env = make_tools()
+    result, events = await _run(tools)
+    assert result["status"] == "success"
+    assert len(result["images"]) == 1 and len(env.stored) == 1
+    assert comfy.submitted[0]["520:517"]["inputs"]["seed"] == 5
+    assert comfy.submitted[0]["520:518"]["inputs"]["value"] == 64
+    assert events[-1]["type"] == "files"
+
+
+async def test_execution_error(make_tools: Callable[[], Any], comfy: FakeComfy) -> None:
+    comfy.script = [
+        {
+            "type": "execution_error",
+            "data": {"prompt_id": "pid-1", "exception_message": "boom"},
+        }
+    ]
+    tools, _ = make_tools()
+    with pytest.raises(RuntimeError, match="boom"):
+        await _run(tools)
+
+
+async def test_ws_drop_retry_declined(
+    make_tools: Callable[[], Any], comfy: FakeComfy
+) -> None:
+    comfy.drop_ws = True
+    tools, _ = make_tools()
+    asked: list[dict[str, Any]] = []
+
+    async def call(e: dict[str, Any]) -> bool:
+        asked.append(e)
+        return False
+
+    with pytest.raises(ConnectionError):
+        await _run(tools, __event_call__=call)
+    assert len(asked) == 1
+
+
+async def test_ws_drop_retry_accepted(
+    make_tools: Callable[[], Any], comfy: FakeComfy
+) -> None:
+    comfy.drop_ws = True
+    tools, _ = make_tools()
+
+    async def call(e: dict[str, Any]) -> bool:
+        comfy.drop_ws = False
+        return True
+
+    result, _ = await _run(tools, __event_call__=call)
+    assert result["status"] == "success"
+
+
+async def test_cancel_interrupts_backend(
+    make_tools: Callable[[], Any], comfy: FakeComfy
+) -> None:
+    comfy.hang = True
+    tools, _ = make_tools()
+    task = asyncio.create_task(_run(tools))
+    while not comfy.submitted:
+        await asyncio.sleep(0.05)
+    await asyncio.sleep(0.1)
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+    assert comfy.interrupts == [{"prompt_id": "pid-1"}]