test: consolidate and fix session runtime tests
This commit is contained in:
Peyton DeNiro
+267
-40
@@ -5,6 +5,17 @@ from unittest.mock import MagicMock, patch, call
|
||||
from comfyui_to_python.runtime_session import WorkflowSessionRuntime
|
||||
|
||||
|
||||
class StubNode:
|
||||
FUNCTION = "execute"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"value": ("STRING",)}}
|
||||
|
||||
def execute(self, value):
|
||||
return (f"result:{value}",)
|
||||
|
||||
|
||||
class TestWorkflowSessionRuntimeInit(unittest.TestCase):
|
||||
"""Tests for WorkflowSessionRuntime initialization."""
|
||||
|
||||
@@ -100,22 +111,6 @@ class TestWorkflowSessionRuntimeLifecycle(unittest.TestCase):
|
||||
self.assertEqual(mock_do_close.call_count, 1)
|
||||
|
||||
|
||||
class TestWorkflowSessionRuntimeRunCount(unittest.TestCase):
|
||||
"""Tests for run count tracking."""
|
||||
|
||||
def _make_runtime(self):
|
||||
return WorkflowSessionRuntime()
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
|
||||
def test_run_increments_run_count(self, mock_do_run):
|
||||
runtime = self._make_runtime()
|
||||
mock_do_run.return_value = {"result": [1]}
|
||||
runtime.run()
|
||||
self.assertEqual(runtime.run_count, 1)
|
||||
runtime.run()
|
||||
self.assertEqual(runtime.run_count, 2)
|
||||
|
||||
|
||||
class TestWorkflowSessionRuntimeExceptionSafety(unittest.TestCase):
|
||||
"""Tests that exceptions during run() do not corrupt session state."""
|
||||
|
||||
@@ -123,21 +118,14 @@ class TestWorkflowSessionRuntimeExceptionSafety(unittest.TestCase):
|
||||
return WorkflowSessionRuntime()
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
|
||||
def test_exception_does_not_clear_bootstrapped(self, mock_do_run):
|
||||
def test_exception_preserves_state_flags(self, mock_do_run):
|
||||
runtime = self._make_runtime()
|
||||
runtime.bootstrapped = True
|
||||
mock_do_run.side_effect = RuntimeError("simulated failure")
|
||||
with self.assertRaises(RuntimeError):
|
||||
runtime.run()
|
||||
self.assertTrue(runtime.bootstrapped)
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
|
||||
def test_exception_preserves_custom_nodes_initialized(self, mock_do_run):
|
||||
runtime = self._make_runtime()
|
||||
runtime.custom_nodes_initialized = True
|
||||
mock_do_run.side_effect = RuntimeError("simulated failure")
|
||||
with self.assertRaises(RuntimeError):
|
||||
runtime.run()
|
||||
self.assertTrue(runtime.bootstrapped)
|
||||
self.assertTrue(runtime.custom_nodes_initialized)
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
|
||||
@@ -148,6 +136,44 @@ class TestWorkflowSessionRuntimeExceptionSafety(unittest.TestCase):
|
||||
runtime.run()
|
||||
self.assertEqual(runtime.run_count, 0)
|
||||
|
||||
def test_run_count_increments_after_successful_run(self):
|
||||
runtime = self._make_runtime()
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
runtime.node_instances = {"StubNode": StubNode()}
|
||||
runtime._node_classes = {"StubNode": StubNode}
|
||||
|
||||
runtime.run()
|
||||
|
||||
self.assertEqual(runtime.run_count, 1)
|
||||
|
||||
def test_run_count_resets_after_reset_every_n_runs(self):
|
||||
runtime = WorkflowSessionRuntime(reset_every_n_runs=2)
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
|
||||
runtime.run()
|
||||
self.assertEqual(runtime.run_count, 1)
|
||||
runtime.run()
|
||||
self.assertEqual(runtime.run_count, 0)
|
||||
self.assertFalse(runtime.bootstrapped)
|
||||
self.assertFalse(runtime.custom_nodes_initialized)
|
||||
self.assertEqual(runtime.node_instances, {})
|
||||
|
||||
|
||||
class TestWorkflowSessionRuntimeAlreadyClosed(unittest.TestCase):
|
||||
"""Tests for behavior after close()."""
|
||||
@@ -168,24 +194,41 @@ class TestWorkflowSessionRuntimeNodeInstances(unittest.TestCase):
|
||||
def _make_runtime(self):
|
||||
return WorkflowSessionRuntime()
|
||||
|
||||
def test_ensure_node_instances_caches_instances(self):
|
||||
def test_ensure_node_instances_creates_and_caches_instances(self):
|
||||
runtime = self._make_runtime()
|
||||
node_class = MagicMock()
|
||||
runtime._ensure_node_instances({"TestNode": node_class})
|
||||
node_class.assert_called_once()
|
||||
self.assertIn("TestNode", runtime.node_instances)
|
||||
|
||||
def test_ensure_node_instances_creates_instance(self):
|
||||
runtime = self._make_runtime()
|
||||
node_class = MagicMock()
|
||||
runtime._ensure_node_instances({"TestNode": node_class})
|
||||
node_class.assert_called_once()
|
||||
|
||||
def test_ensure_node_instances_reuses_cached(self):
|
||||
runtime = self._make_runtime()
|
||||
node_class = MagicMock()
|
||||
runtime._ensure_node_instances({"TestNode": node_class})
|
||||
runtime._ensure_node_instances({"TestNode": node_class})
|
||||
node_class.assert_called_once()
|
||||
def test_node_instances_are_cached_across_runs(self):
|
||||
runtime = WorkflowSessionRuntime()
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
original_stub_init = StubNode.__init__
|
||||
init_calls = []
|
||||
|
||||
def tracking_init(self, *args, **kwargs):
|
||||
init_calls.append(1)
|
||||
original_stub_init(self)
|
||||
|
||||
StubNode.__init__ = tracking_init
|
||||
try:
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = {"StubNode": StubNode}
|
||||
|
||||
runtime.run()
|
||||
runtime.run()
|
||||
|
||||
self.assertEqual(len(init_calls), 1)
|
||||
finally:
|
||||
StubNode.__init__ = original_stub_init
|
||||
|
||||
|
||||
class TestWorkflowSessionRuntimeClearRuntimeCache(unittest.TestCase):
|
||||
@@ -225,17 +268,201 @@ class TestWorkflowSessionRuntimeDoClose(unittest.TestCase):
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.cleanup_comfyui_runtime")
|
||||
@patch("comfyui_to_python.runtime_session.gc")
|
||||
def test_do_close_calls_cleanup_with_unload(self, mock_gc, mock_cleanup):
|
||||
def test_do_close_calls_cleanup_and_gc(self, mock_gc, mock_cleanup):
|
||||
runtime = self._make_runtime()
|
||||
runtime._do_close(unload_models=True)
|
||||
mock_cleanup.assert_called_once_with(unload_models=True)
|
||||
mock_gc.collect.assert_called_once()
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.cleanup_comfyui_runtime")
|
||||
@patch("comfyui_to_python.runtime_session.gc")
|
||||
def test_do_close_calls_cleanup_without_unload(self, mock_gc, mock_cleanup):
|
||||
mock_cleanup.reset_mock()
|
||||
mock_gc.reset_mock()
|
||||
runtime = self._make_runtime()
|
||||
runtime._do_close(unload_models=False)
|
||||
mock_cleanup.assert_called_once_with(unload_models=False)
|
||||
mock_gc.collect.assert_called_once()
|
||||
|
||||
|
||||
class TestWorkflowSessionRuntimeRun(unittest.TestCase):
|
||||
"""Tests for WorkflowSessionRuntime.run() workflow execution."""
|
||||
|
||||
def _make_runtime(self):
|
||||
return WorkflowSessionRuntime()
|
||||
|
||||
def test_run_executes_workflow_nodes(self):
|
||||
runtime = self._make_runtime()
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
runtime.node_instances = {"StubNode": StubNode()}
|
||||
runtime._node_classes = {"StubNode": StubNode}
|
||||
|
||||
result = runtime.run()
|
||||
|
||||
self.assertIn("1", result)
|
||||
self.assertEqual(result["1"], ["result:test"])
|
||||
|
||||
@patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
|
||||
)
|
||||
def test_run_with_no_workflow_data_returns_none(self, mock_ensure):
|
||||
runtime = self._make_runtime()
|
||||
result = runtime.run()
|
||||
self.assertIsNone(result)
|
||||
|
||||
@patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
|
||||
)
|
||||
def test_run_with_no_node_mappings_returns_none(self, mock_ensure):
|
||||
runtime = self._make_runtime()
|
||||
runtime._workflow_data = {"1": {"class_type": "Test", "inputs": {}}}
|
||||
result = runtime.run()
|
||||
self.assertIsNone(result)
|
||||
|
||||
@patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
|
||||
)
|
||||
def test_run_node_not_in_mappings_is_skipped(self, mock_ensure):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "UnknownNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = self._make_runtime()
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
|
||||
result = runtime.run()
|
||||
|
||||
self.assertEqual(result, {})
|
||||
|
||||
def test_run_node_returns_tuple_is_converted_to_list(self):
|
||||
runtime = self._make_runtime()
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "multi"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
runtime.node_instances = {"StubNode": StubNode()}
|
||||
runtime._node_classes = {"StubNode": StubNode}
|
||||
|
||||
result = runtime.run()
|
||||
|
||||
self.assertIsInstance(result["1"], list)
|
||||
self.assertEqual(result["1"][0], "result:multi")
|
||||
|
||||
def test_run_persists_parameters_across_calls(self):
|
||||
runtime = self._make_runtime()
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "persist"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
|
||||
with patch.object(
|
||||
runtime, "_do_run", return_value={"1": ["result:persist"]}
|
||||
) as mock_do_run:
|
||||
runtime.run(
|
||||
workflow_data=workflow_data, node_class_mappings=node_mappings
|
||||
)
|
||||
runtime.run()
|
||||
|
||||
self.assertEqual(mock_do_run.call_count, 2)
|
||||
|
||||
def test_run_does_not_corrupt_session_on_exception(self):
|
||||
runtime = self._make_runtime()
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
mock_instance = MagicMock()
|
||||
mock_instance.execute.side_effect = ValueError("boom")
|
||||
runtime.node_instances = {"StubNode": mock_instance}
|
||||
runtime._node_classes = {"StubNode": StubNode}
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
runtime.run()
|
||||
|
||||
runtime.node_instances = {"StubNode": StubNode()}
|
||||
runtime._node_classes = {"StubNode": StubNode}
|
||||
result = runtime.run()
|
||||
self.assertEqual(result["1"], ["result:test"])
|
||||
|
||||
@patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
|
||||
)
|
||||
def test_session_policy_does_not_clear_cache(self, mock_ensure):
|
||||
runtime = WorkflowSessionRuntime(cleanup_policy="session")
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
|
||||
with patch.object(runtime, "clear_runtime_cache") as mock_clear:
|
||||
runtime.run()
|
||||
mock_clear.assert_not_called()
|
||||
|
||||
@patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
|
||||
)
|
||||
def test_per_run_policy_clears_cache_after_each_run(self, mock_ensure):
|
||||
runtime = WorkflowSessionRuntime(cleanup_policy="per_run")
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
|
||||
with patch.object(runtime, "clear_runtime_cache") as mock_clear:
|
||||
runtime.run()
|
||||
mock_clear.assert_called_once()
|
||||
|
||||
@patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
|
||||
)
|
||||
def test_manual_policy_never_clears_cache(self, mock_ensure):
|
||||
runtime = WorkflowSessionRuntime(cleanup_policy="manual")
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
|
||||
with patch.object(runtime, "clear_runtime_cache") as mock_clear:
|
||||
runtime.run()
|
||||
runtime.run()
|
||||
mock_clear.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,107 +1,75 @@
|
||||
import threading
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from comfyui_to_python import WorkflowSession
|
||||
from comfyui_to_python.runtime_session import WorkflowSession
|
||||
|
||||
|
||||
class TestWorkflowSessionInit(unittest.TestCase):
|
||||
"""Tests for WorkflowSession initialization."""
|
||||
|
||||
def test_init_default_cleanup_policy(self):
|
||||
def test_init_creates_runtime(self):
|
||||
session = WorkflowSession()
|
||||
self.assertEqual(session._runtime._cleanup_policy, "session")
|
||||
self.assertIsNotNone(session._runtime)
|
||||
|
||||
def test_init_accepts_per_run_policy(self):
|
||||
def test_init_passes_cleanup_policy(self):
|
||||
session = WorkflowSession(cleanup_policy="per_run")
|
||||
self.assertEqual(session._runtime._cleanup_policy, "per_run")
|
||||
|
||||
def test_init_accepts_manual_policy(self):
|
||||
session = WorkflowSession(cleanup_policy="manual")
|
||||
self.assertEqual(session._runtime._cleanup_policy, "manual")
|
||||
|
||||
def test_init_accepts_reset_every_n_runs(self):
|
||||
def test_init_passes_reset_every_n_runs(self):
|
||||
session = WorkflowSession(reset_every_n_runs=5)
|
||||
self.assertEqual(session._runtime._reset_every_n_runs, 5)
|
||||
|
||||
def test_init_wraps_runtime(self):
|
||||
session = WorkflowSession()
|
||||
self.assertIsInstance(session._runtime, object)
|
||||
self.assertTrue(hasattr(session._runtime, "run"))
|
||||
self.assertTrue(hasattr(session._runtime, "close"))
|
||||
self.assertTrue(hasattr(session._runtime, "clear_runtime_cache"))
|
||||
|
||||
class TestWorkflowSessionDelegation(unittest.TestCase):
|
||||
"""Tests for WorkflowSession public API delegation to internal runtime."""
|
||||
|
||||
class TestWorkflowSessionRun(unittest.TestCase):
|
||||
"""Tests for WorkflowSession.run() delegation."""
|
||||
def _make_session(self):
|
||||
return WorkflowSession()
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
|
||||
def test_run_delegates_to_runtime(self, mock_do_run):
|
||||
mock_do_run.return_value = {"result": [1, 2, 3]}
|
||||
session = WorkflowSession()
|
||||
result = session.run()
|
||||
self.assertEqual(result, {"result": [1, 2, 3]})
|
||||
def test_run_delegates_to_runtime(self):
|
||||
session = self._make_session()
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": MagicMock()}
|
||||
session._runtime.node_instances = {"StubNode": MagicMock()}
|
||||
session._runtime._node_classes = {"StubNode": MagicMock()}
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
|
||||
def test_run_with_workflow_data(self, mock_do_run):
|
||||
workflow_data = {"1": {"class_type": "TestNode", "inputs": {"x": 1}}}
|
||||
node_mappings = {"TestNode": MagicMock()}
|
||||
session = WorkflowSession()
|
||||
session.run(workflow_data=workflow_data, node_class_mappings=node_mappings)
|
||||
mock_do_run.assert_called_once()
|
||||
with patch.object(
|
||||
session._runtime, "run", return_value={"1": ["result"]}
|
||||
) as mock_run:
|
||||
session.run(workflow_data=workflow_data, node_class_mappings=node_mappings)
|
||||
mock_run.assert_called_once_with(
|
||||
workflow_data=workflow_data,
|
||||
node_class_mappings=node_mappings,
|
||||
extra_pnginfo=None,
|
||||
)
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
|
||||
def test_run_with_extra_pnginfo(self, mock_do_run):
|
||||
extra_pnginfo = {"workflow": {"metadata": "test"}}
|
||||
session = WorkflowSession()
|
||||
session.run(extra_pnginfo=extra_pnginfo)
|
||||
mock_do_run.assert_called_once()
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
|
||||
def test_run_passes_all_parameters(self, mock_do_run):
|
||||
workflow_data = {"1": {"class_type": "TestNode", "inputs": {}}}
|
||||
node_mappings = {"TestNode": MagicMock()}
|
||||
extra_pnginfo = {"workflow": {}}
|
||||
session = WorkflowSession()
|
||||
session.run(
|
||||
workflow_data=workflow_data,
|
||||
node_class_mappings=node_mappings,
|
||||
extra_pnginfo=extra_pnginfo,
|
||||
)
|
||||
mock_do_run.assert_called_once()
|
||||
|
||||
|
||||
class TestWorkflowSessionClose(unittest.TestCase):
|
||||
"""Tests for WorkflowSession.close() delegation."""
|
||||
def test_clear_runtime_cache_delegates_to_runtime(self):
|
||||
session = self._make_session()
|
||||
with patch.object(
|
||||
session._runtime, "clear_runtime_cache"
|
||||
) as mock_clear:
|
||||
session.clear_runtime_cache()
|
||||
mock_clear.assert_called_once()
|
||||
|
||||
def test_close_delegates_to_runtime(self):
|
||||
session = WorkflowSession()
|
||||
with patch.object(session._runtime, "close") as mock_close:
|
||||
session = self._make_session()
|
||||
with patch.object(
|
||||
session._runtime, "close"
|
||||
) as mock_close:
|
||||
session.close(unload_models=True)
|
||||
mock_close.assert_called_once_with(unload_models=True)
|
||||
|
||||
def test_close_delegates_with_unload_false(self):
|
||||
session = WorkflowSession()
|
||||
with patch.object(session._runtime, "close") as mock_close:
|
||||
session.close(unload_models=False)
|
||||
mock_close.assert_called_once_with(unload_models=False)
|
||||
|
||||
def test_run_after_close_raises(self):
|
||||
session = WorkflowSession()
|
||||
def test_run_raises_after_close(self):
|
||||
session = self._make_session()
|
||||
session.close()
|
||||
with self.assertRaises(RuntimeError):
|
||||
session.run()
|
||||
|
||||
|
||||
class TestWorkflowSessionClearCache(unittest.TestCase):
|
||||
"""Tests for WorkflowSession.clear_runtime_cache() delegation."""
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime.clear_runtime_cache")
|
||||
def test_clear_cache_delegates_to_runtime(self, mock_clear):
|
||||
session = WorkflowSession()
|
||||
session.clear_runtime_cache()
|
||||
mock_clear.assert_called_once()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -1,270 +0,0 @@
|
||||
import json
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from comfyui_to_python.runtime_session import WorkflowSessionRuntime
|
||||
|
||||
|
||||
class StubNode:
|
||||
"""Stub node class for testing."""
|
||||
|
||||
FUNCTION = "execute"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"value": ("STRING",)}}
|
||||
|
||||
def execute(self, value):
|
||||
return (f"result:{value}",)
|
||||
|
||||
|
||||
class TestWorkflowSessionRuntimeDoRun(unittest.TestCase):
|
||||
"""Tests for WorkflowSessionRuntime._do_run() workflow execution."""
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_close")
|
||||
def test_run_executes_workflow_nodes(self, mock_do_close):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
runtime.node_instances = {"StubNode": StubNode()}
|
||||
runtime._node_classes = {"StubNode": StubNode}
|
||||
|
||||
result = runtime.run()
|
||||
|
||||
self.assertIn("1", result)
|
||||
self.assertEqual(result["1"], ["result:test"])
|
||||
|
||||
@patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
|
||||
)
|
||||
def test_run_with_no_workflow_data_returns_none(self, mock_ensure):
|
||||
runtime = WorkflowSessionRuntime()
|
||||
result = runtime.run()
|
||||
self.assertIsNone(result)
|
||||
|
||||
@patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
|
||||
)
|
||||
def test_run_with_no_node_mappings_returns_none(self, mock_ensure):
|
||||
runtime = WorkflowSessionRuntime()
|
||||
runtime._workflow_data = {"1": {"class_type": "Test", "inputs": {}}}
|
||||
result = runtime.run()
|
||||
self.assertIsNone(result)
|
||||
|
||||
@patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
|
||||
)
|
||||
def test_run_node_not_in_mappings_is_skipped(self, mock_unknown):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "UnknownNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
|
||||
result = runtime.run()
|
||||
|
||||
# Unknown node is skipped, but the call still returns empty dict
|
||||
# since no node instances were created
|
||||
self.assertEqual(result, {})
|
||||
|
||||
def test_run_node_returns_tuple_is_converted_to_list(self):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "multi"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
runtime.node_instances = {"StubNode": StubNode()}
|
||||
runtime._node_classes = {"StubNode": StubNode}
|
||||
|
||||
result = runtime.run()
|
||||
|
||||
self.assertIsInstance(result["1"], list)
|
||||
self.assertEqual(result["1"][0], "result:multi")
|
||||
|
||||
def test_run_preserves_extra_pnginfo(self):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
extra_pnginfo = {"workflow": {"test": True}}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
runtime._extra_pnginfo = extra_pnginfo
|
||||
runtime.node_instances = {"StubNode": StubNode()}
|
||||
runtime._node_classes = {"StubNode": StubNode}
|
||||
|
||||
result = runtime.run()
|
||||
|
||||
self.assertIn("1", result)
|
||||
|
||||
def test_run_persists_parameters_across_calls(self):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "persist"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
|
||||
with patch.object(
|
||||
runtime, "_do_run", return_value={"1": ["result:persist"]}
|
||||
) as mock_do_run:
|
||||
runtime.run(
|
||||
workflow_data=workflow_data, node_class_mappings=node_mappings
|
||||
)
|
||||
runtime.run()
|
||||
|
||||
self.assertEqual(mock_do_run.call_count, 2)
|
||||
|
||||
def test_run_does_not_corrupt_session_on_exception(self):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
mock_instance = MagicMock()
|
||||
mock_instance.execute.side_effect = ValueError("boom")
|
||||
runtime.node_instances = {"StubNode": mock_instance}
|
||||
runtime._node_classes = {"StubNode": StubNode}
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
runtime.run()
|
||||
|
||||
# Session should still be usable
|
||||
node_mappings["StubNode"] = StubNode
|
||||
runtime.node_instances = {"StubNode": StubNode()}
|
||||
runtime._node_classes = {"StubNode": StubNode}
|
||||
result = runtime.run()
|
||||
self.assertEqual(result["1"], ["result:test"])
|
||||
|
||||
@patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
|
||||
)
|
||||
def test_run_increments_run_count(self, mock_ensure):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
|
||||
runtime.run()
|
||||
self.assertEqual(runtime.run_count, 1)
|
||||
runtime.run()
|
||||
self.assertEqual(runtime.run_count, 2)
|
||||
|
||||
@patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
|
||||
)
|
||||
def test_session_policy_does_not_clear_cache(self, mock_ensure):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime(cleanup_policy="session")
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
|
||||
with patch.object(runtime, "clear_runtime_cache") as mock_clear:
|
||||
runtime.run()
|
||||
mock_clear.assert_not_called()
|
||||
|
||||
@patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
|
||||
)
|
||||
def test_per_run_policy_clears_cache_after_each_run(self, mock_ensure):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime(cleanup_policy="per_run")
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
|
||||
with patch.object(runtime, "clear_runtime_cache") as mock_clear:
|
||||
runtime.run()
|
||||
mock_clear.assert_called_once()
|
||||
|
||||
@patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
|
||||
)
|
||||
def test_manual_policy_never_clears_cache(self, mock_ensure):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime(cleanup_policy="manual")
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
|
||||
with patch.object(runtime, "clear_runtime_cache") as mock_clear:
|
||||
runtime.run()
|
||||
runtime.run()
|
||||
mock_clear.assert_not_called()
|
||||
|
||||
def test_node_instances_are_cached(self):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
|
||||
with patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run",
|
||||
return_value={"1": ["result:test"]},
|
||||
):
|
||||
runtime.run()
|
||||
runtime.run()
|
||||
|
||||
# Node class should only be instantiated once
|
||||
self.assertIn("StubNode", runtime.node_instances)
|
||||
self.assertEqual(len(runtime.node_instances), 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user