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
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
|
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"}]
@pytest.fixture
def fast_idle(monkeypatch: pytest.MonkeyPatch) -> None:
"""Shrink the 1s WebSocket idle timeout so stall detection takes ms, not seconds."""
import comfier
real_wait_for = asyncio.wait_for
async def wait_for(aw: Any, timeout: float | None = None) -> Any:
return await real_wait_for(aw, 0.01 if timeout == 1.0 else timeout)
monkeypatch.setattr(comfier.asyncio, "wait_for", wait_for)
async def test_stalled_backend_lost_prompt(
make_tools: Callable[[], Any], comfy: FakeComfy, fast_idle: None
) -> None:
comfy.hang = True
comfy.queue_has_prompt = False # not in /queue, and /history is empty
tools, _ = make_tools()
with pytest.raises(ConnectionError, match="lost prompt pid-1"):
await asyncio.wait_for(_run(tools), timeout=5)
async def test_stalled_backend_still_queued_keeps_waiting(
make_tools: Callable[[], Any], comfy: FakeComfy, fast_idle: None
) -> None:
comfy.hang = True # silent, but /queue still lists the prompt
tools, _ = make_tools()
task = asyncio.create_task(_run(tools))
while not comfy.submitted:
await asyncio.sleep(0.01)
await asyncio.sleep(0.3) # many idle ticks and liveness checks pass
assert not task.done()
await comfy.release()
result, _ = await asyncio.wait_for(task, timeout=5)
assert result["status"] == "success"
|
