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.
110 lines
2.8 KiB
Python
110 lines
2.8 KiB
Python
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()
|