Peyton DeNiro
08ad731038
Add session-mode code generation that produces a reusable WorkflowSession class instead of a simple oneshot script. The session mode generates code with an in-class run/close lifecycle for warm ComfyUI reuse. Also refactors render.py to extract shared code sections into common methods, reducing duplication between oneshot and session renderers. Adds WorkflowSessionRuntime with cleanup policies (per_run, session, manual) and reset_every_n_runs support, plus 72 new tests covering session rendering, runtime session, and session export pipeline.
108 lines
4.1 KiB
Python
108 lines
4.1 KiB
Python
import threading
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from comfyui_to_python import WorkflowSession
|
|
|
|
|
|
class TestWorkflowSessionInit(unittest.TestCase):
|
|
"""Tests for WorkflowSession initialization."""
|
|
|
|
def test_init_default_cleanup_policy(self):
|
|
session = WorkflowSession()
|
|
self.assertEqual(session._runtime._cleanup_policy, "session")
|
|
|
|
def test_init_accepts_per_run_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):
|
|
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 TestWorkflowSessionRun(unittest.TestCase):
|
|
"""Tests for WorkflowSession.run() delegation."""
|
|
|
|
@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]})
|
|
|
|
@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()
|
|
|
|
@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_close_delegates_to_runtime(self):
|
|
session = WorkflowSession()
|
|
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()
|
|
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()
|