diff --git a/tests/test_runtime_session.py b/tests/test_runtime_session.py index 9e1b286..d670c3b 100644 --- a/tests/test_runtime_session.py +++ b/tests/test_runtime_session.py @@ -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__": diff --git a/tests/test_workflow_session.py b/tests/test_workflow_session.py index 25bbb4a..2c9ffb3 100644 --- a/tests/test_workflow_session.py +++ b/tests/test_workflow_session.py @@ -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() diff --git a/tests/test_workflow_session_runtime_execution.py b/tests/test_workflow_session_runtime_execution.py deleted file mode 100644 index 5d2ca72..0000000 --- a/tests/test_workflow_session_runtime_execution.py +++ /dev/null @@ -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()