256 lines
9.1 KiB
Python
256 lines
9.1 KiB
Python
"""Offline tests for the CLI WebSocket protocol handling.
|
|
|
|
These mock out websocket and HTTP so we can verify the client behavior
|
|
against representative ComfyUI server message sequences without needing a
|
|
live server.
|
|
"""
|
|
import json
|
|
import sys
|
|
import types
|
|
|
|
import pytest
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
@pytest.fixture
|
|
def hymt2_cli(monkeypatch):
|
|
"""Import cli/hymt2_cli with a stubbed `websocket` module."""
|
|
# Build a stub websocket module that the test can drive per-test.
|
|
stub = types.ModuleType("websocket")
|
|
stub.WebSocketTimeoutException = TimeoutError
|
|
stub.WebSocketConnectionClosedException = ConnectionError
|
|
# Default WebSocket class — tests override with their own
|
|
stub.WebSocket = lambda: _NullWS()
|
|
monkeypatch.setitem(sys.modules, "websocket", stub)
|
|
|
|
# Add cli/ to sys.path BEFORE importing
|
|
import os.path as osp
|
|
cli_dir = osp.abspath(osp.join(osp.dirname(__file__), "..", "cli"))
|
|
if cli_dir not in sys.path:
|
|
sys.path.insert(0, cli_dir)
|
|
|
|
# Force re-import each test
|
|
sys.modules.pop("hymt2_cli", None)
|
|
import hymt2_cli as mod
|
|
return mod, stub
|
|
|
|
|
|
class _NullWS:
|
|
def connect(self, *a, **k): pass
|
|
def settimeout(self, t): pass
|
|
def recv(self): raise RuntimeError("unexpected recv")
|
|
def close(self): pass
|
|
|
|
|
|
class _ScriptedWS:
|
|
"""Plays back a fixed sequence of WS messages."""
|
|
def __init__(self, seq):
|
|
self.seq = list(seq)
|
|
self.i = 0
|
|
self.closed = False
|
|
|
|
def connect(self, *a, **k): pass
|
|
def settimeout(self, t): pass
|
|
|
|
def recv(self):
|
|
if self.i >= len(self.seq):
|
|
raise RuntimeError("seq exhausted (test should have completed earlier)")
|
|
msg = self.seq[self.i]
|
|
self.i += 1
|
|
return json.dumps(msg)
|
|
|
|
def close(self): self.closed = True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests
|
|
# ---------------------------------------------------------------------------
|
|
class TestCachedPromptRace:
|
|
"""v0.1.2 critical regression test: cached prompt completes between
|
|
submit() and the first ws.recv(). The client must detect this via the
|
|
/history poll right after submit(), not block on recv()."""
|
|
|
|
def test_cache_hit_returns_immediately(self, hymt2_cli):
|
|
mod, ws_stub = hymt2_cli
|
|
|
|
# WS only sends a status/queue=0 message and would hang on next recv
|
|
ws_stub.WebSocket = lambda: _ScriptedWS([
|
|
{"type": "status",
|
|
"data": {"status": {"exec_info": {"queue_remaining": 0}}}},
|
|
])
|
|
|
|
client = mod.ComfyClient(server="x:1", verbose=False)
|
|
client.submit = lambda wf: "P-CACHE"
|
|
# /history reports completed immediately after submit
|
|
client.history = lambda pid: {
|
|
"status": {"completed": True, "status_str": "success"},
|
|
"outputs": {"preview": {"text": ["cached result"]}},
|
|
}
|
|
|
|
result = client.run({}, timeout_s=5)
|
|
text = mod.extract_text_from_history(result, "preview")
|
|
assert text == "cached result"
|
|
|
|
|
|
class TestColdLoadPath:
|
|
"""When the prompt isn't cached, the WS sees the full event stream
|
|
ending with execution_success."""
|
|
|
|
def test_full_event_stream(self, hymt2_cli):
|
|
mod, ws_stub = hymt2_cli
|
|
|
|
ws_stub.WebSocket = lambda: _ScriptedWS([
|
|
{"type": "status",
|
|
"data": {"status": {"exec_info": {"queue_remaining": 1}}}},
|
|
{"type": "execution_start",
|
|
"data": {"prompt_id": "P-COLD"}},
|
|
{"type": "executing",
|
|
"data": {"node": "loader", "prompt_id": "P-COLD"}},
|
|
{"type": "progress",
|
|
"data": {"value": 1, "max": 4, "prompt_id": "P-COLD"}},
|
|
{"type": "progress",
|
|
"data": {"value": 4, "max": 4, "prompt_id": "P-COLD"}},
|
|
{"type": "executing",
|
|
"data": {"node": "translate", "prompt_id": "P-COLD"}},
|
|
{"type": "execution_success",
|
|
"data": {"prompt_id": "P-COLD"}},
|
|
])
|
|
|
|
# First /history (right after submit) shows still running, second
|
|
# (after WS loop) shows completed.
|
|
history_state = {"calls": 0}
|
|
|
|
def fake_history(pid):
|
|
history_state["calls"] += 1
|
|
if history_state["calls"] == 1:
|
|
return {"status": {"completed": False}}
|
|
return {
|
|
"status": {"completed": True, "status_str": "success"},
|
|
"outputs": {"preview": {"text": ["cold result"]}},
|
|
}
|
|
|
|
client = mod.ComfyClient(server="x:1", verbose=False)
|
|
client.submit = lambda wf: "P-COLD"
|
|
client.history = fake_history
|
|
|
|
result = client.run({}, timeout_s=10)
|
|
text = mod.extract_text_from_history(result, "preview")
|
|
assert text == "cold result"
|
|
|
|
|
|
class TestLegacyPath:
|
|
"""The pre-execution_success protocol (executing/node=null as the queue
|
|
drain signal) must still work."""
|
|
|
|
def test_legacy_executing_null(self, hymt2_cli):
|
|
mod, ws_stub = hymt2_cli
|
|
|
|
ws_stub.WebSocket = lambda: _ScriptedWS([
|
|
{"type": "execution_start",
|
|
"data": {"prompt_id": "P-LEGACY"}},
|
|
{"type": "executing",
|
|
"data": {"node": "translate", "prompt_id": "P-LEGACY"}},
|
|
{"type": "executing",
|
|
"data": {"node": None, "prompt_id": "P-LEGACY"}},
|
|
])
|
|
|
|
client = mod.ComfyClient(server="x:1", verbose=False)
|
|
client.submit = lambda wf: "P-LEGACY"
|
|
client.history = lambda pid: {
|
|
"status": {"completed": True, "status_str": "success"},
|
|
"outputs": {"preview": {"text": ["legacy result"]}},
|
|
}
|
|
|
|
result = client.run({}, timeout_s=10)
|
|
assert mod.extract_text_from_history(result, "preview") == "legacy result"
|
|
|
|
|
|
class TestErrorPath:
|
|
"""execution_error for our prompt must raise."""
|
|
|
|
def test_execution_error_raises(self, hymt2_cli):
|
|
mod, ws_stub = hymt2_cli
|
|
|
|
ws_stub.WebSocket = lambda: _ScriptedWS([
|
|
{"type": "execution_start",
|
|
"data": {"prompt_id": "P-ERR"}},
|
|
{"type": "execution_error",
|
|
"data": {"prompt_id": "P-ERR",
|
|
"exception_message": "CUDA OOM"}},
|
|
])
|
|
|
|
client = mod.ComfyClient(server="x:1", verbose=False)
|
|
client.submit = lambda wf: "P-ERR"
|
|
client.history = lambda pid: {}
|
|
|
|
with pytest.raises(RuntimeError) as exc_info:
|
|
client.run({}, timeout_s=10)
|
|
assert "execution_error" in str(exc_info.value)
|
|
assert "CUDA OOM" in str(exc_info.value)
|
|
|
|
|
|
class TestPromptIdIsolation:
|
|
"""Concurrent prompts must not cross-trigger completion."""
|
|
|
|
def test_ignores_unrelated_prompt_completion(self, hymt2_cli):
|
|
mod, ws_stub = hymt2_cli
|
|
|
|
ws_stub.WebSocket = lambda: _ScriptedWS([
|
|
{"type": "execution_success",
|
|
"data": {"prompt_id": "SOMEONE_ELSE"}},
|
|
{"type": "execution_success",
|
|
"data": {"prompt_id": "P-MINE"}},
|
|
])
|
|
|
|
client = mod.ComfyClient(server="x:1", verbose=False)
|
|
client.submit = lambda wf: "P-MINE"
|
|
client.history = lambda pid: {
|
|
"status": {"completed": True},
|
|
"outputs": {"preview": {"text": ["correct"]}},
|
|
}
|
|
|
|
result = client.run({}, timeout_s=10)
|
|
assert mod.extract_text_from_history(result, "preview") == "correct"
|
|
|
|
|
|
class TestExtractText:
|
|
"""extract_text_from_history handles list and string output shapes."""
|
|
|
|
def test_text_as_list(self):
|
|
# Re-import without ws stub for plain function tests
|
|
import os.path as osp, sys as _sys
|
|
cli_dir = osp.abspath(osp.join(osp.dirname(__file__), "..", "cli"))
|
|
if cli_dir not in _sys.path:
|
|
_sys.path.insert(0, cli_dir)
|
|
# Need to stub websocket again because module-level import requires it
|
|
if "websocket" not in _sys.modules:
|
|
stub = types.ModuleType("websocket")
|
|
stub.WebSocketTimeoutException = TimeoutError
|
|
stub.WebSocketConnectionClosedException = ConnectionError
|
|
stub.WebSocket = lambda: _NullWS()
|
|
_sys.modules["websocket"] = stub
|
|
_sys.modules.pop("hymt2_cli", None)
|
|
import hymt2_cli as mod
|
|
|
|
h = {"outputs": {"preview": {"text": ["Good morning."]}}}
|
|
assert mod.extract_text_from_history(h, "preview") == "Good morning."
|
|
|
|
def test_missing_node(self):
|
|
import os.path as osp, sys as _sys
|
|
cli_dir = osp.abspath(osp.join(osp.dirname(__file__), "..", "cli"))
|
|
if cli_dir not in _sys.path:
|
|
_sys.path.insert(0, cli_dir)
|
|
if "websocket" not in _sys.modules:
|
|
stub = types.ModuleType("websocket")
|
|
stub.WebSocketTimeoutException = TimeoutError
|
|
stub.WebSocketConnectionClosedException = ConnectionError
|
|
stub.WebSocket = lambda: _NullWS()
|
|
_sys.modules["websocket"] = stub
|
|
_sys.modules.pop("hymt2_cli", None)
|
|
import hymt2_cli as mod
|
|
|
|
h = {"outputs": {"other": {"text": ["x"]}}}
|
|
assert mod.extract_text_from_history(h, "preview") == ""
|