"""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") == ""