From 08ad73103809a6eac1bf321c45924613567440d3 Mon Sep 17 00:00:00 2001 From: Peyton DeNiro <25550995+pydn@users.noreply.github.com> Date: Sun, 19 Apr 2026 12:06:46 -0500 Subject: [PATCH] feat: add session execution mode with WorkflowSession and WorkflowSessionRuntime 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. --- comfyui_to_python/__init__.py | 4 + comfyui_to_python/app.py | 5 +- comfyui_to_python/generator/model.py | 5 +- comfyui_to_python/generator/planner.py | 3 + comfyui_to_python/generator/render.py | 288 ++++++++++++++---- comfyui_to_python/runtime_session.py | 205 +++++++++++++ conftest.py | 6 + tests/conftest.py | 16 + tests/test_export_session_mode.py | 109 +++++++ tests/test_runtime_session.py | 242 +++++++++++++++ tests/test_session_renderer.py | 273 +++++++++++++++++ tests/test_workflow_session.py | 107 +++++++ ...test_workflow_session_runtime_execution.py | 270 ++++++++++++++++ 13 files changed, 1472 insertions(+), 61 deletions(-) create mode 100644 comfyui_to_python/runtime_session.py create mode 100644 conftest.py create mode 100644 tests/conftest.py create mode 100644 tests/test_export_session_mode.py create mode 100644 tests/test_runtime_session.py create mode 100644 tests/test_session_renderer.py create mode 100644 tests/test_workflow_session.py create mode 100644 tests/test_workflow_session_runtime_execution.py diff --git a/comfyui_to_python/__init__.py b/comfyui_to_python/__init__.py index 18aafa1..5c35c93 100644 --- a/comfyui_to_python/__init__.py +++ b/comfyui_to_python/__init__.py @@ -3,6 +3,7 @@ from typing import TextIO from .app import ExportApplication from .cli import DEFAULT_INPUT_FILE, DEFAULT_OUTPUT_FILE, DEFAULT_QUEUE_SIZE, main from .node_runtime import get_node_class_mappings, import_custom_nodes +from .runtime_session import WorkflowSession class ComfyUItoPython: @@ -17,6 +18,7 @@ class ComfyUItoPython: queue_size: int = 1, node_class_mappings: dict | None = None, needs_init_custom_nodes: bool = False, + execution_mode: str = "oneshot", ): self._app = ExportApplication( workflow=workflow, @@ -26,6 +28,7 @@ class ComfyUItoPython: queue_size=queue_size, node_class_mappings=node_class_mappings, needs_init_custom_nodes=needs_init_custom_nodes, + execution_mode=execution_mode, node_mapping_loader=get_node_class_mappings, custom_node_importer=import_custom_nodes, ) @@ -48,6 +51,7 @@ def run( __all__ = [ "ComfyUItoPython", + "WorkflowSession", "run", "main", "get_node_class_mappings", diff --git a/comfyui_to_python/app.py b/comfyui_to_python/app.py index 334d06f..4d584a8 100644 --- a/comfyui_to_python/app.py +++ b/comfyui_to_python/app.py @@ -23,6 +23,7 @@ class ExportApplication: needs_init_custom_nodes: bool = False, node_mapping_loader=None, custom_node_importer=None, + execution_mode: str = "oneshot", ): if input_file and workflow: raise ValueError("Can't provide both input_file and workflow") @@ -44,6 +45,7 @@ class ExportApplication: else self.node_mapping_loader() ) self.needs_init_custom_nodes = needs_init_custom_nodes + self.execution_mode = execution_mode self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings) def execute(self) -> None: @@ -69,7 +71,8 @@ class ExportApplication: data, metadata_workflow_data, queue_size=self.queue_size, + execution_mode=self.execution_mode, ) - generated_code = WorkflowRenderer().render(plan) + generated_code = WorkflowRenderer(execution_mode=self.execution_mode).render(plan) write_python_output(self.output_file, generated_code) print(f"Code successfully generated and written to {self.output_file}") diff --git a/comfyui_to_python/generator/model.py b/comfyui_to_python/generator/model.py index ffe5e89..235d114 100644 --- a/comfyui_to_python/generator/model.py +++ b/comfyui_to_python/generator/model.py @@ -1,4 +1,5 @@ -from dataclasses import dataclass +from dataclasses import dataclass, field +from typing import Literal @dataclass(frozen=True) @@ -10,3 +11,5 @@ class GenerationPlan: metadata_workflow_data: dict | None queue_size: int custom_nodes: bool + execution_mode: Literal["oneshot", "session"] = field(default="oneshot") + executed_variables: dict[str, str] = field(default_factory=dict) diff --git a/comfyui_to_python/generator/planner.py b/comfyui_to_python/generator/planner.py index cfc5090..d9f4e52 100644 --- a/comfyui_to_python/generator/planner.py +++ b/comfyui_to_python/generator/planner.py @@ -37,6 +37,7 @@ class WorkflowPlanner: workflow_data: dict, metadata_workflow_data: dict | None = None, queue_size: int = 10, + execution_mode: str = "oneshot", ) -> GenerationPlan: import_statements = {"nodes": {"NODE_CLASS_MAPPINGS"}} executed_variables = {} @@ -132,6 +133,8 @@ class WorkflowPlanner: metadata_workflow_data=metadata_workflow_data, queue_size=queue_size, custom_nodes=custom_nodes, + execution_mode=execution_mode, + executed_variables=executed_variables, ) def create_function_call_code( diff --git a/comfyui_to_python/generator/render.py b/comfyui_to_python/generator/render.py index 1acd4a3..0b95a89 100644 --- a/comfyui_to_python/generator/render.py +++ b/comfyui_to_python/generator/render.py @@ -1,4 +1,5 @@ import inspect +import threading from pprint import pformat from typing import Any @@ -20,15 +21,25 @@ from .model import GenerationPlan class WorkflowRenderer: """Render a generation plan into the final standalone Python source.""" - def render(self, plan: GenerationPlan) -> str: - workflow_literal = self.format_python_literal(plan.workflow_data) - if plan.metadata_workflow_data is None: - extra_pnginfo_literal = "None" - else: - extra_pnginfo_literal = self.format_python_literal( - {"workflow": plan.metadata_workflow_data} - ) + def __init__(self, execution_mode: str = "oneshot"): + self.execution_mode = execution_mode + def render(self, plan: GenerationPlan) -> str: + if self.execution_mode == "session": + return self._render_session_mode(plan) + return self._render_oneshot_mode(plan) + + # ── shared sections ────────────────────────────────────────────── + + @staticmethod + def _build_entrypoint_section() -> list[str]: + return [ + "# Entrypoint", + 'if __name__ == "__main__":', + " main()", + ] + + def _build_imports_section(self, plan: GenerationPlan) -> list[str]: func_strings = [] for func in [ get_value_at_index, @@ -52,16 +63,19 @@ class WorkflowRenderer: if plan.custom_nodes: static_imports.append(f"\n{inspect.getsource(import_custom_nodes)}\n") - custom_nodes_call = "import_custom_nodes()" + + return static_imports + + def _build_workflow_section(self, plan: GenerationPlan) -> list[str]: + workflow_literal = self.format_python_literal(plan.workflow_data) + if plan.metadata_workflow_data is None: + extra_pnginfo_literal = "None" else: - custom_nodes_call = None + extra_pnginfo_literal = self.format_python_literal( + {"workflow": plan.metadata_workflow_data} + ) - imports_code = [] - for module_name in sorted(plan.import_statements.keys()): - class_names = ", ".join(sorted(plan.import_statements[module_name])) - imports_code.append(f"from {module_name} import {class_names}") - - workflow_section = [ + return [ "# Workflow data", "def build_workflow() -> dict[str, Any]:", f" return {workflow_literal}", @@ -74,52 +88,16 @@ class WorkflowRenderer: "extra_pnginfo = build_extra_pnginfo()", ] - execution_section = [ - "# Workflow execution", - "def main(unload_models: bool | None = None):", - " bootstrap_comfyui_runtime()", - " add_extra_model_paths()", - ] - if custom_nodes_call: - execution_section.append(f" {custom_nodes_call}") - if imports_code: - execution_section.extend(["", " # Node imports"]) - execution_section.extend(f" {line}" for line in imports_code) - execution_section.extend( - [ - "", - " import torch", - "", - " try:", - " with torch.inference_mode():", - ] - ) - execution_section.extend( - self.build_function_body( - plan.special_functions_code, "pass", indentation=" " - ).splitlines() - ) - execution_section.append(f" for q in range({plan.queue_size}):") - execution_section.extend( - self.build_function_body( - plan.loop_code, "pass", indentation=" " - ).splitlines() - ) - execution_section.extend( - [ - " finally:", - " cleanup_comfyui_runtime(unload_models=unload_models)", - ] - ) + # ── oneshot renderer ───────────────────────────────────────────── - entrypoint_section = [ - "# Entrypoint", - 'if __name__ == "__main__":', - " main()", - ] + def _render_oneshot_mode(self, plan: GenerationPlan) -> str: + imports_section = self._build_imports_section(plan) + workflow_section = self._build_workflow_section(plan) + execution_section = self._build_execution_section(plan) + entrypoint_section = self._build_entrypoint_section() final_code = "\n".join( - static_imports + imports_section + [""] + workflow_section + [""] @@ -129,6 +107,198 @@ class WorkflowRenderer: ) return black.format_str(final_code, mode=black.Mode()) + def _build_execution_section(self, plan: GenerationPlan) -> list[str]: + imports_code = self._build_node_imports(plan.import_statements) + + lines = [ + "# Workflow execution", + "def main(unload_models: bool | None = None):", + " bootstrap_comfyui_runtime()", + " add_extra_model_paths()", + ] + if plan.custom_nodes: + lines.append(" import_custom_nodes()") + if imports_code: + lines.extend(["", " # Node imports"]) + lines.extend(f" {line}" for line in imports_code) + lines.extend( + [ + "", + " import torch", + "", + " try:", + " with torch.inference_mode():", + ] + ) + lines.extend( + self.build_function_body( + plan.special_functions_code, "pass", indentation=" " + ).splitlines() + ) + lines.append(f" for q in range({plan.queue_size}):") + lines.extend( + self.build_function_body( + plan.loop_code, "pass", indentation=" " + ).splitlines() + ) + lines.extend( + [ + " finally:", + " cleanup_comfyui_runtime(unload_models=unload_models)", + ] + ) + return lines + + @staticmethod + def _build_node_imports( + import_statements: dict[str, set[str]], + ) -> list[str]: + imports_code = [] + for module_name in sorted(import_statements.keys()): + class_names = ", ".join(sorted(import_statements[module_name])) + imports_code.append(f"from {module_name} import {class_names}") + return imports_code + + # ── session renderer ───────────────────────────────────────────── + + def _render_session_mode(self, plan: GenerationPlan) -> str: + imports_section = self._build_imports_section(plan) + workflow_section = self._build_workflow_section(plan) + session_class = self._build_session_class(plan) + main_wrapper = self._build_main_wrapper(plan) + entrypoint_section = self._build_entrypoint_section() + + final_code = "\n".join( + imports_section + + [""] + + session_class + + [""] + + workflow_section + + [""] + + main_wrapper + + [""] + + entrypoint_section + ) + return black.format_str(final_code, mode=black.Mode()) + + def _build_session_class(self, plan: GenerationPlan) -> list[str]: + lines = [ + "# WorkflowSession class", + "class WorkflowSession:", + ' """A reusable warm-session wrapper for generated ComfyUI workflows."""', + "", + ' def __init__(self, cleanup_policy: str = "per_run", reset_every_n_runs: int | None = None):', + ' """Initialize the session.', + "", + ' Args:', + ' cleanup_policy: One of "per_run", "session", or "manual".', + ' reset_every_n_runs: If set, soft-reset every N runs.', + ' """', + " self._bootstrapped = False", + " self._custom_nodes_initialized = False", + " self._node_instances = {}", + " self._lock = threading.Lock()", + " self._closed = False", + " self._cleanup_policy = cleanup_policy", + " self._reset_every_n_runs = reset_every_n_runs", + " self._run_count = 0", + f" self._queue_size = {plan.queue_size}", + "", + " def run(self) -> dict[str, Any] | None:", + ' """Run the workflow and return the output (or None)."""', + " with self._lock:", + " if self._closed:", + " raise RuntimeError('Session is closed')", + "", + " if not self._bootstrapped:", + " self._bootstrapped = True", + " bootstrap_comfyui_runtime()", + "", + " if self._cleanup_policy == 'session' and not self._custom_nodes_initialized:", + " self._custom_nodes_initialized = True", + " import_custom_nodes()", + "", + " prompt = json.loads(json.dumps(build_workflow()))", + " extra_pnginfo = build_extra_pnginfo()", + "", + " import torch", + " try:", + " with torch.inference_mode():", + ] + + # Add special functions body (inside inference_mode) + special_body = self.build_function_body( + plan.special_functions_code, "pass", indentation=" " + ) + lines.extend(special_body.splitlines()) + + lines.append(" for q in range(self._queue_size):") + + # Add loop code (node instantiations + calls) + loop_body = self.build_function_body( + plan.loop_code, "pass", indentation=" " + ) + lines.extend(loop_body.splitlines()) + + lines.extend( + [ + " finally:", + " if self._cleanup_policy == 'per_run':", + " cleanup_comfyui_runtime(unload_models=True)", + "", + ] + ) + + # Build outputs collection: outputs = {node_id: var_name, ...} + executed_vars = plan.executed_variables + if executed_vars: + outputs_init = " outputs = {}" + outputs_assigns = [] + for node_id, var_name in executed_vars.items(): + outputs_assigns.append(f" outputs[{node_id!r}] = {var_name}") + run_increment = " self._run_count += 1" + outputs_return = " return outputs" + lines.append(outputs_init) + lines.extend(outputs_assigns) + lines.append(run_increment) + lines.append(outputs_return) + else: + run_increment = " self._run_count += 1" + outputs_return = " return None" + lines.append(run_increment) + lines.append(outputs_return) + + lines.extend( + [ + "", + " def close(self, unload_models: bool = True) -> None:", + ' """Close the session and optionally unload models."""', + " with self._lock:", + " if self._closed:", + " return", + " self._closed = True", + " cleanup_comfyui_runtime(unload_models=unload_models)", + " import gc", + " gc.collect()", + ] + ) + + return lines + + def _build_main_wrapper(self, plan: GenerationPlan) -> list[str]: + return [ + "# Entry point", + "def main(unload_models: bool | None = None):", + ' """Backward-compatible entry point using a short-lived WorkflowSession."""', + ' session = WorkflowSession(cleanup_policy="per_run")', + " try:", + " session.run()", + " finally:", + " session.close(unload_models=unload_models)", + ] + + # ── helpers ────────────────────────────────────────────────────── + @staticmethod def format_python_literal(value: Any) -> str: return pformat(value, sort_dicts=False) diff --git a/comfyui_to_python/runtime_session.py b/comfyui_to_python/runtime_session.py new file mode 100644 index 0000000..6c8661b --- /dev/null +++ b/comfyui_to_python/runtime_session.py @@ -0,0 +1,205 @@ +import gc +import json +import threading +from typing import Any, Literal + +from comfyui_to_python.node_runtime import ( + bootstrap_comfyui_runtime, + cleanup_comfyui_runtime, + import_custom_nodes, +) +from comfyui_to_python.generator.model import GenerationPlan + + +_CLEANUP_POLICIES = {"per_run", "session", "manual"} + + +class WorkflowSessionRuntime: + """Internal runtime session for warm reuse of ComfyUI state.""" + + def __init__( + self, + cleanup_policy: Literal["per_run", "session", "manual"] = "session", + reset_every_n_runs: int | None = None, + ): + if cleanup_policy not in _CLEANUP_POLICIES: + raise ValueError( + f"cleanup_policy must be one of {sorted(_CLEANUP_POLICIES)}, " + f"got {cleanup_policy!r}" + ) + + self._cleanup_policy: Literal["per_run", "session", "manual"] = cleanup_policy + self._reset_every_n_runs: int | None = reset_every_n_runs + + self.bootstrapped: bool = False + self.custom_nodes_initialized: bool = False + self.node_instances: dict[str, Any] = {} + self._node_classes: dict[str, Any] = {} + self.run_count: int = 0 + self._closed: bool = False + self._lock: threading.Lock = threading.Lock() + self._workflow_data: dict | None = None + self._node_class_mappings: dict | None = None + self._extra_pnginfo: dict | None = None + + def _ensure_bootstrapped(self) -> None: + if self.bootstrapped: + return + bootstrap_comfyui_runtime() + self.bootstrapped = True + + def _ensure_custom_nodes_initialized(self) -> None: + if self.custom_nodes_initialized: + return + import_custom_nodes() + self.custom_nodes_initialized = True + + def _ensure_node_instances(self, node_classes: dict) -> None: + for class_type, node_class in node_classes.items(): + if class_type in self.node_instances: + continue + self.node_instances[class_type] = node_class() + self._node_classes[class_type] = node_class + + def clear_runtime_cache(self) -> None: + if self._cleanup_policy == "session": + return + if self._cleanup_policy == "per_run": + cleanup_comfyui_runtime(unload_models=True) + + def close(self, unload_models: bool = True) -> None: + if self._closed: + return + self._do_close(unload_models=unload_models) + self._closed = True + + def _do_close(self, unload_models: bool = True) -> None: + cleanup_comfyui_runtime(unload_models=unload_models) + gc.collect() + + def _do_run(self) -> Any: + if not self._workflow_data or not self._node_class_mappings: + return None + + prompt = json.loads(json.dumps(self._workflow_data)) + extra_pnginfo = self._extra_pnginfo if hasattr(self, "_extra_pnginfo") else None + + try: + import torch + + inference_ctx = torch.inference_mode + except ImportError: + + def inference_ctx(): + class _DummyCtx: + def __enter__(self): + pass + + def __exit__(self, *args): + pass + + return _DummyCtx() + + with inference_ctx(): + outputs = {} + for node_id, node in prompt.items(): + class_type = node.get("class_type", "") + inputs = node.get("inputs", {}) + if class_type not in self.node_instances: + continue + node_instance = self.node_instances[class_type] + node_class = self._node_classes.get(class_type) + func_name = getattr(node_class, "FUNCTION", "execute") + func = getattr(node_instance, func_name, None) + if func is None: + continue + args = {} + for k, v in inputs.items(): + args[k] = v + result = func(**args) + if isinstance(result, tuple) or isinstance(result, list): + outputs[node_id] = list(result) + else: + outputs[node_id] = result + + return outputs + + def run( + self, + workflow_data: dict | None = None, + node_class_mappings: dict | None = None, + extra_pnginfo: dict | None = None, + ) -> Any: + with self._lock: + if self._closed: + raise RuntimeError( + "Cannot run() on a closed WorkflowSessionRuntime" + ) + + self._workflow_data = workflow_data or self._workflow_data + self._node_class_mappings = ( + node_class_mappings or self._node_class_mappings + ) + self._extra_pnginfo = extra_pnginfo or self._extra_pnginfo + + if self._node_class_mappings: + self._ensure_node_instances(self._node_class_mappings) + + try: + result = self._do_run() + except Exception: + raise + else: + if self._cleanup_policy == "per_run": + self.clear_runtime_cache() + self.run_count += 1 + + if ( + self._reset_every_n_runs + and self.run_count % self._reset_every_n_runs == 0 + ): + self._do_close(unload_models=False) + self._closed = False + self.bootstrapped = False + self.custom_nodes_initialized = False + self.node_instances.clear() + self._node_classes.clear() + self.run_count = 0 + + return result + + +class WorkflowSession: + """A reusable warm-session wrapper for generated ComfyUI workflows. + + This is the public API class that wraps WorkflowSessionRuntime. + It delegates all method calls to the internal runtime instance. + """ + + def __init__( + self, + cleanup_policy: Literal["per_run", "session", "manual"] = "session", + reset_every_n_runs: int | None = None, + ): + self._runtime = WorkflowSessionRuntime( + cleanup_policy=cleanup_policy, + reset_every_n_runs=reset_every_n_runs, + ) + + def run( + self, + workflow_data: dict | None = None, + node_class_mappings: dict | None = None, + extra_pnginfo: dict | None = None, + ) -> Any: + return self._runtime.run( + workflow_data=workflow_data, + node_class_mappings=node_class_mappings, + extra_pnginfo=extra_pnginfo, + ) + + def clear_runtime_cache(self) -> None: + self._runtime.clear_runtime_cache() + + def close(self, unload_models: bool = True) -> None: + self._runtime.close(unload_models=unload_models) diff --git a/conftest.py b/conftest.py new file mode 100644 index 0000000..615cb8f --- /dev/null +++ b/conftest.py @@ -0,0 +1,6 @@ +"""Mock the ComfyUI 'server' module so the root __init__.py can be imported by pytest.""" + +import sys +from unittest.mock import MagicMock + +sys.modules["server"] = MagicMock() diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..5812095 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,16 @@ +"""Pytest configuration to handle the ComfyUI extension __init__.py at repo root.""" + +import sys +from pathlib import Path + +import pytest + + +@pytest.fixture(autouse=True) +def _setup_test_path(): + """Ensure the package is importable during tests.""" + # Add the repo root to sys.path so `from comfyui_to_python import ...` works + repo_root = Path(__file__).parent.parent.resolve() + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + yield diff --git a/tests/test_export_session_mode.py b/tests/test_export_session_mode.py new file mode 100644 index 0000000..c0c5eed --- /dev/null +++ b/tests/test_export_session_mode.py @@ -0,0 +1,109 @@ +import json +import unittest +from io import StringIO + +from comfyui_to_python import ComfyUItoPython + + +class DummyNode: + CATEGORY = "test" + FUNCTION = "execute" + + @classmethod + def INPUT_TYPES(cls): + return {"required": {"value": ("STRING",)}} + + def execute(self, value): + return (f"result:{value}",) + + +class ExportSessionModeTest(unittest.TestCase): + """Tests for session mode export pipeline.""" + + def test_comfyui_to_python_passes_execution_mode(self): + workflow = { + "1": { + "class_type": "DummyNode", + "inputs": {"value": "test"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"DummyNode": DummyNode}, + execution_mode="session", + ) + + generated = output.getvalue() + + self.assertIn("class WorkflowSession:", generated) + self.assertIn("def run(self)", generated) + self.assertIn("def close(self, unload_models:", generated) + + def test_session_mode_includes_backward_compat_main(self): + workflow = { + "1": { + "class_type": "DummyNode", + "inputs": {"value": "test"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"DummyNode": DummyNode}, + execution_mode="session", + ) + + generated = output.getvalue() + + self.assertIn("def main(", generated) + self.assertIn("WorkflowSession(", generated) + self.assertIn('cleanup_policy="per_run"', generated) + self.assertIn("if __name__ == \"__main__\":", generated) + + def test_oneshot_mode_excludes_session_class(self): + workflow = { + "1": { + "class_type": "DummyNode", + "inputs": {"value": "test"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"DummyNode": DummyNode}, + execution_mode="oneshot", + ) + + generated = output.getvalue() + + self.assertNotIn("class WorkflowSession:", generated) + + def test_default_mode_is_oneshot(self): + workflow = { + "1": { + "class_type": "DummyNode", + "inputs": {"value": "test"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"DummyNode": DummyNode}, + ) + + generated = output.getvalue() + + self.assertNotIn("class WorkflowSession:", generated) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_runtime_session.py b/tests/test_runtime_session.py new file mode 100644 index 0000000..9e1b286 --- /dev/null +++ b/tests/test_runtime_session.py @@ -0,0 +1,242 @@ +import threading +import unittest +from unittest.mock import MagicMock, patch, call + +from comfyui_to_python.runtime_session import WorkflowSessionRuntime + + +class TestWorkflowSessionRuntimeInit(unittest.TestCase): + """Tests for WorkflowSessionRuntime initialization.""" + + def test_init_default_cleanup_policy_is_session(self): + runtime = WorkflowSessionRuntime() + self.assertEqual(runtime._cleanup_policy, "session") + + def test_init_accepts_per_run_policy(self): + runtime = WorkflowSessionRuntime(cleanup_policy="per_run") + self.assertEqual(runtime._cleanup_policy, "per_run") + + def test_init_accepts_manual_policy(self): + runtime = WorkflowSessionRuntime(cleanup_policy="manual") + self.assertEqual(runtime._cleanup_policy, "manual") + + def test_init_rejects_invalid_policy(self): + with self.assertRaises(ValueError): + WorkflowSessionRuntime(cleanup_policy="invalid") + + def test_init_sets_default_reset_every_n_runs_none(self): + runtime = WorkflowSessionRuntime() + self.assertIsNone(runtime._reset_every_n_runs) + + def test_init_accepts_reset_every_n_runs(self): + runtime = WorkflowSessionRuntime(reset_every_n_runs=5) + self.assertEqual(runtime._reset_every_n_runs, 5) + + def test_init_has_lock(self): + runtime = WorkflowSessionRuntime() + self.assertIsInstance(runtime._lock, type(threading.Lock())) + + def test_init_bootstrapped_false(self): + runtime = WorkflowSessionRuntime() + self.assertFalse(runtime.bootstrapped) + + def test_init_custom_nodes_initialized_false(self): + runtime = WorkflowSessionRuntime() + self.assertFalse(runtime.custom_nodes_initialized) + + def test_init_run_count_zero(self): + runtime = WorkflowSessionRuntime() + self.assertEqual(runtime.run_count, 0) + + def test_init_closed_false(self): + runtime = WorkflowSessionRuntime() + self.assertFalse(runtime._closed) + + +class TestWorkflowSessionRuntimeLifecycle(unittest.TestCase): + """Tests for WorkflowSessionRuntime lifecycle management.""" + + def _make_runtime(self): + return WorkflowSessionRuntime() + + @patch("comfyui_to_python.runtime_session.bootstrap_comfyui_runtime") + def test_ensure_bootstrapped_calls_bootstrap_once(self, mock_bootstrap): + runtime = self._make_runtime() + runtime._ensure_bootstrapped() + mock_bootstrap.assert_called_once() + self.assertTrue(runtime.bootstrapped) + + @patch("comfyui_to_python.runtime_session.bootstrap_comfyui_runtime") + def test_ensure_bootstrapped_skips_if_already_bootstrapped(self, mock_bootstrap): + runtime = self._make_runtime() + runtime._ensure_bootstrapped() + runtime._ensure_bootstrapped() + self.assertEqual(mock_bootstrap.call_count, 1) + + @patch("comfyui_to_python.runtime_session.import_custom_nodes") + def test_ensure_custom_nodes_init_calls_import_once(self, mock_import): + runtime = self._make_runtime() + runtime._ensure_custom_nodes_initialized() + mock_import.assert_called_once() + self.assertTrue(runtime.custom_nodes_initialized) + + @patch("comfyui_to_python.runtime_session.import_custom_nodes") + def test_ensure_custom_nodes_init_skips_if_already_initialized(self, mock_import): + runtime = self._make_runtime() + runtime._ensure_custom_nodes_initialized() + runtime._ensure_custom_nodes_initialized() + self.assertEqual(mock_import.call_count, 1) + + def test_close_sets_closed_flag(self): + runtime = self._make_runtime() + runtime.close(unload_models=True) + self.assertTrue(runtime._closed) + + def test_close_is_idempotent(self): + runtime = self._make_runtime() + with patch.object(runtime, "_do_close") as mock_do_close: + runtime.close(unload_models=True) + runtime.close(unload_models=True) + 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.""" + + def _make_runtime(self): + return WorkflowSessionRuntime() + + @patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run") + def test_exception_does_not_clear_bootstrapped(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.custom_nodes_initialized) + + @patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run") + def test_run_count_not_incremented_on_exception(self, mock_do_run): + runtime = self._make_runtime() + mock_do_run.side_effect = RuntimeError("simulated failure") + with self.assertRaises(RuntimeError): + runtime.run() + self.assertEqual(runtime.run_count, 0) + + +class TestWorkflowSessionRuntimeAlreadyClosed(unittest.TestCase): + """Tests for behavior after close().""" + + def _make_runtime(self): + return WorkflowSessionRuntime() + + def test_run_after_close_raises(self): + runtime = self._make_runtime() + runtime.close() + with self.assertRaises(RuntimeError): + runtime.run() + + +class TestWorkflowSessionRuntimeNodeInstances(unittest.TestCase): + """Tests for cached node instance management.""" + + def _make_runtime(self): + return WorkflowSessionRuntime() + + def test_ensure_node_instances_caches_instances(self): + runtime = self._make_runtime() + node_class = MagicMock() + runtime._ensure_node_instances({"TestNode": node_class}) + 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() + + +class TestWorkflowSessionRuntimeClearRuntimeCache(unittest.TestCase): + """Tests for clear_runtime_cache behavior.""" + + def _make_runtime(self): + return WorkflowSessionRuntime() + + @patch("comfyui_to_python.runtime_session.cleanup_comfyui_runtime") + def test_clear_runtime_cache_session_policy_skips_unload(self, mock_cleanup): + runtime = WorkflowSessionRuntime(cleanup_policy="session") + runtime.bootstrapped = True + runtime.clear_runtime_cache() + # session policy should NOT call unload_all_models + mock_cleanup.assert_not_called() + + @patch("comfyui_to_python.runtime_session.cleanup_comfyui_runtime") + def test_clear_runtime_cache_per_run_policy_calls_full_cleanup(self, mock_cleanup): + runtime = WorkflowSessionRuntime(cleanup_policy="per_run") + runtime.bootstrapped = True + runtime.clear_runtime_cache() + mock_cleanup.assert_called_once_with(unload_models=True) + + @patch("comfyui_to_python.runtime_session.cleanup_comfyui_runtime") + def test_clear_runtime_cache_manual_policy_no_cleanup(self, mock_cleanup): + runtime = WorkflowSessionRuntime(cleanup_policy="manual") + runtime.bootstrapped = True + runtime.clear_runtime_cache() + mock_cleanup.assert_not_called() + + +class TestWorkflowSessionRuntimeDoClose(unittest.TestCase): + """Tests for _do_close internal method.""" + + def _make_runtime(self): + return WorkflowSessionRuntime() + + @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): + runtime = self._make_runtime() + runtime._do_close(unload_models=True) + mock_cleanup.assert_called_once_with(unload_models=True) + + @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): + runtime = self._make_runtime() + runtime._do_close(unload_models=False) + mock_cleanup.assert_called_once_with(unload_models=False) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_session_renderer.py b/tests/test_session_renderer.py new file mode 100644 index 0000000..14acd06 --- /dev/null +++ b/tests/test_session_renderer.py @@ -0,0 +1,273 @@ +import json +import unittest +from io import StringIO +from unittest.mock import patch + +from comfyui_to_python import ComfyUItoPython + + +class LoadImage: + CATEGORY = "image" + FUNCTION = "load_image" + + @classmethod + def INPUT_TYPES(cls): + return {"required": {"image": ("STRING",)}} + + def load_image(self, image): + return (image,) + + +class SessionRendererTest(unittest.TestCase): + """Tests for session mode code generation in the renderer.""" + + def test_session_mode_generates_workflow_session_class(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": {"image": "example.png"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"LoadImage": LoadImage}, + execution_mode="session", + ) + + generated = output.getvalue() + self.assertIn("class WorkflowSession", generated) + + def test_session_mode_generates_main_wrapper(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": {"image": "example.png"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"LoadImage": LoadImage}, + execution_mode="session", + ) + + generated = output.getvalue() + self.assertIn("def main(unload_models: bool | None = None)", generated) + self.assertIn("WorkflowSession(", generated) + + def test_session_mode_main_creates_session_with_per_run_policy(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": {"image": "example.png"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"LoadImage": LoadImage}, + execution_mode="session", + ) + + generated = output.getvalue() + self.assertIn('cleanup_policy="per_run"', generated) + + def test_session_mode_main_has_try_finally_close(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": {"image": "example.png"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"LoadImage": LoadImage}, + execution_mode="session", + ) + + generated = output.getvalue() + self.assertIn("try:", generated) + self.assertIn("session.run()", generated) + self.assertIn("finally:", generated) + self.assertIn("session.close(", generated) + + def test_session_mode_generates_run_method(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": {"image": "example.png"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"LoadImage": LoadImage}, + execution_mode="session", + ) + + generated = output.getvalue() + self.assertIn("def run(self", generated) + + def test_session_mode_generates_close_method(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": {"image": "example.png"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"LoadImage": LoadImage}, + execution_mode="session", + ) + + generated = output.getvalue() + self.assertIn("def close(self, unload_models: bool = True)", generated) + + def test_session_mode_oneshot_generates_same_code(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": {"image": "example.png"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"LoadImage": LoadImage}, + execution_mode="oneshot", + ) + + generated = output.getvalue() + self.assertNotIn("class WorkflowSession", generated) + self.assertIn("def main(unload_models: bool | None = None)", generated) + self.assertIn("bootstrap_comfyui_runtime()", generated) + self.assertIn("cleanup_comfyui_runtime(unload_models=unload_models)", generated) + + def test_session_mode_oneshot_default(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": {"image": "example.png"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"LoadImage": LoadImage}, + ) + + generated = output.getvalue() + self.assertNotIn("class WorkflowSession", generated) + self.assertIn("def main(unload_models: bool | None = None)", generated) + + def test_session_mode_script_is_executable(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": {"image": "example.png"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"LoadImage": LoadImage}, + execution_mode="session", + ) + + generated = output.getvalue() + + # Should not raise — checks that generated code has valid syntax + compile(generated, "", "exec") + + # main() should be callable + globals_dict = {"__name__": "generated_workflow_module"} + with patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime"): + exec(generated, globals_dict) + self.assertIn("WorkflowSession", globals_dict) + self.assertIn("main", globals_dict) + self.assertTrue(callable(globals_dict["main"])) + + def test_session_mode_generates_workflow_literal(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": {"image": "example.png"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"LoadImage": LoadImage}, + execution_mode="session", + ) + + generated = output.getvalue() + self.assertIn("def build_workflow()", generated) + self.assertIn('return', generated) + self.assertIn('"class_type": "LoadImage"', generated) + + def test_session_mode_generates_bootstrap_helper(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": {"image": "example.png"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"LoadImage": LoadImage}, + execution_mode="session", + ) + + generated = output.getvalue() + self.assertIn("def bootstrap_comfyui_runtime()", generated) + + def test_session_mode_generates_cleanup_helper(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": {"image": "example.png"}, + } + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={"LoadImage": LoadImage}, + execution_mode="session", + ) + + generated = output.getvalue() + self.assertIn("def cleanup_comfyui_runtime(", generated) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_workflow_session.py b/tests/test_workflow_session.py new file mode 100644 index 0000000..25bbb4a --- /dev/null +++ b/tests/test_workflow_session.py @@ -0,0 +1,107 @@ +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() diff --git a/tests/test_workflow_session_runtime_execution.py b/tests/test_workflow_session_runtime_execution.py new file mode 100644 index 0000000..5d2ca72 --- /dev/null +++ b/tests/test_workflow_session_runtime_execution.py @@ -0,0 +1,270 @@ +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()