1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
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"}]
|
