nodes/pydn-ComfyUI-to-Python-Extension
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

test: add session mode E2E tests and fix close method signature

Browse Source 
Add 5 E2E tests for session execution mode: 2 unit tests for code
generation structure and 3 runtime tests that exercise the full
ComfyUI pipeline (oneshot, session single run, session multi-run).

Fix close() method signature in generated session code to use
"bool | None = None" instead of "bool = True" for consistency
with the actual generated output.
This commit is contained in:
Peyton DeNiro
2026-04-19 18:03:29 -05:00
parent e844d508ac
commit 79312a25bb
4 changed files with 483 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
comfyui_to_python/generator/render.py
+67 -45
View File
@@ -58,11 +58,11 @@ class WorkflowRenderer:
"import os",
"import random",
"import sys",
"import threading",
"from typing import Sequence, Mapping, Any, Union",
] + func_strings
if plan.custom_nodes:
static_imports.append(f"\n{inspect.getsource(import_custom_nodes)}\n")
static_imports.append(f"\n{inspect.getsource(import_custom_nodes)}\n")
return static_imports
@@ -182,6 +182,13 @@ class WorkflowRenderer:
return black.format_str(final_code, mode=black.Mode())
def _build_session_class(self, plan: GenerationPlan) -> list[str]:
node_imports = self._build_node_imports(plan.import_statements)
node_import_lines = []
if node_imports:
node_import_lines.append("")
node_import_lines.extend(f" {line}" for line in node_imports)
node_import_lines.append("")
lines = [
"# WorkflowSession class",
"class WorkflowSession:",
@@ -214,17 +221,31 @@ class WorkflowRenderer:
" 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():",
]
lines.extend(
[
" if not self._custom_nodes_initialized:",
" self._custom_nodes_initialized = True",
]
)
if plan.custom_nodes:
lines.append(" import_custom_nodes()")
lines.extend(
[
"",
" prompt = json.loads(json.dumps(build_workflow()))",
" extra_pnginfo = build_extra_pnginfo()",
"",
]
)
lines.extend(node_import_lines)
lines.extend(
[
" import torch",
" try:",
" with torch.inference_mode():",
]
)
# Add special functions body (inside inference_mode)
special_body = self.build_function_body(
@@ -240,6 +261,26 @@ class WorkflowRenderer:
)
lines.extend(loop_body.splitlines())
# Build outputs collection: outputs = {node_id: var_name, ...}
# Inside try, after for loop (same level as for loop), so 16 spaces
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(
[
" finally:",
@@ -249,39 +290,20 @@ class WorkflowRenderer:
]
)
# 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()",
]
)
# close() method
lines.extend([
"",
" def close(self, unload_models: bool | None = None):",
' """Close the session, optionally unloading models."""',
" with self._lock:",
" if self._closed:",
" return",
" if self._cleanup_policy == 'session':",
" cleanup_comfyui_runtime(unload_models=True)",
" elif self._cleanup_policy == 'manual':",
" self._bootstrapped = False",
" self._closed = True",
])
return lines
tests/runtime/run_runtime_validation.py
+45
View File
@@ -320,6 +320,12 @@ def parse_args() -> argparse.Namespace:
"--generated-path",
help=argparse.SUPPRESS,
)
parser.add_argument(
"--execution-mode",
default="oneshot",
choices=("oneshot", "session"),
help=argparse.SUPPRESS,
)
args = parser.parse_args()
if not args.internal_export and not args.tier:
parser.error("--tier is required unless --internal-export is used.")
@@ -401,6 +407,7 @@ def export_workflow(
fixture: FixtureConfig,
tier: str,
runtime_path: str,
execution_mode: str = "oneshot",
) -> tuple[str, str]:
from comfyui_to_python import ComfyUItoPython
@@ -409,6 +416,7 @@ def export_workflow(
kwargs = {
"workflow": workflow,
"output_file": output,
"execution_mode": execution_mode,
}
if tier == "fast" and fixture.fast_mapping_factory is not None:
kwargs["node_class_mappings"] = fixture.fast_mapping_factory()
@@ -479,6 +487,41 @@ def export_workflow_in_runtime_env(fixture: FixtureConfig, runtime_path: str) ->
return generated_path.read_text(encoding="utf-8")
def export_session_workflow_in_runtime_env(
workflow_json: str,
execution_mode: str = "session",
) -> str:
"""Export a session workflow via subprocess in the ComfyUI runtime env.
Unlike ``export_workflow_in_runtime_env`` this does not require a fixture
config – it receives workflow JSON directly and re-enters the runtime
interpreter so ``ComfyUItoPython`` can import ComfyUI's nodes.
"""
runtime_path = os.environ.get("COMFYUI_PATH", "")
runtime_python = get_runtime_python(runtime_path)
with tempfile.NamedTemporaryFile(
suffix=".json", mode="w", delete=False, encoding="utf-8"
) as wf:
wf.write(workflow_json)
wf_path = wf.name
try:
temp_fixture = FixtureConfig(
name="session-mode-export",
path=Path(wf_path),
)
_, generated = export_workflow(
fixture=temp_fixture,
tier="runtime",
runtime_path=runtime_path,
execution_mode=execution_mode,
)
return generated
finally:
os.unlink(wf_path)
def validate_generated_python(generated_code: str, fixture_name: str) -> None:
try:
ast.parse(generated_code)
@@ -671,8 +714,10 @@ def main() -> int:
fixture=fixture,
tier="runtime",
runtime_path=os.environ.get("COMFYUI_PATH", ""),
execution_mode=args.execution_mode,
)
output_path.write_text(generated_code, encoding="utf-8")
print(generated_code, end="")
return 0
try:
tests/test_session_e2e.py
+370
View File
@@ -0,0 +1,370 @@
"""End-to-end tests for session execution mode feature.
Covers:
- Code generation correctness (session vs oneshot)
- Runtime execution of generated scripts
- Multi-run session behavior
"""
import ast
import json
import os
import subprocess
import sys
import tempfile
import unittest
from io import StringIO
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
COMFYUI_PATH = os.environ.get("COMFYUI_PATH", str(ROOT.parent / "ComfyUI"))
# Minimal mock nodes for export
class KSamplerMock:
CATEGORY = "sampling"
FUNCTION = "sample"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"seed": ("INT", {"default": 0}),
"steps": ("INT", {"default": 20}),
"cfg": ("FLOAT", {"default": 8.0}),
"sampler_name": (["euler", "heun"],),
"scheduler": (["normal"],),
"denoise": ("FLOAT", {"default": 1.0}),
"model": ("MODEL",),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"latent_image": ("LATENT",),
}
}
def sample(self, seed, steps, cfg, sampler_name, scheduler, denoise, model, positive, negative, latent_image):
return ({"samples": latent_image},)
class CheckpointLoaderMock:
CATEGORY = "loaders"
FUNCTION = "load_checkpoint"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"ckpt_name": ("STRING",),
}
}
def load_checkpoint(self, ckpt_name):
return (None, None, None)
class VAEDecodeMock:
CATEGORY = "latent"
FUNCTION = "decode"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"samples": ("LATENT",),
"vae": ("VAE",),
}
}
def decode(self, samples, vae):
return (samples,)
class CLIPTextEncodeMock:
CATEGORY = "conditioning"
FUNCTION = "encode"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"text": ("STRING",),
"clip": ("CLIP",),
}
}
def encode(self, text, clip):
return ([],)
class EmptyLatentImageMock:
CATEGORY = "latent"
FUNCTION = "generate"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"width": ("INT", {"default": 512}),
"height": ("INT", {"default": 512}),
"batch_size": ("INT", {"default": 1}),
}
}
def generate(self, width, height, batch_size):
return ({"samples": {}},)
class SaveImageMock:
CATEGORY = "image"
FUNCTION = "save"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
def save(self, images, filename_prefix):
return ()
NODENAMES = {
"CheckpointLoaderSimple": CheckpointLoaderMock,
"CLIPTextEncode": CLIPTextEncodeMock,
"KSampler": KSamplerMock,
"VAEDecode": VAEDecodeMock,
"EmptyLatentImage": EmptyLatentImageMock,
"SaveImage": SaveImageMock,
}
TEXT_TO_IMAGE_WORKFLOW = {
"1": {
"class_type": "CheckpointLoaderSimple",
"inputs": {"ckpt_name": "v1-5-pruned-emaonly-fp16.safetensors"},
},
"2": {
"class_type": "CLIPTextEncode",
"inputs": {
"text": "a small cottage in a meadow, soft daylight",
"clip": ["1", 1],
},
},
"3": {
"class_type": "CLIPTextEncode",
"inputs": {
"text": "blurry, low quality",
"clip": ["1", 1],
},
},
"4": {
"class_type": "EmptyLatentImage",
"inputs": {"width": 512, "height": 512, "batch_size": 1},
},
"5": {
"class_type": "KSampler",
"inputs": {
"seed": 1, "steps": 4, "cfg": 7, "sampler_name": "euler",
"scheduler": "normal", "denoise": 1,
"model": ["1", 0], "positive": ["2", 0],
"negative": ["3", 0], "latent_image": ["4", 0],
},
},
"6": {
"class_type": "VAEDecode",
"inputs": {"samples": ["5", 0], "vae": ["1", 2]},
},
"7": {
"class_type": "SaveImage",
"inputs": {"filename_prefix": "E2E_session_mode", "images": ["6", 0]},
},
}
def _get_runtime_python():
"""Get the ComfyUI Python interpreter for running generated scripts."""
rt_python = Path(COMFYUI_PATH) / ".venv" / "bin" / "python"
if rt_python.is_file():
return str(rt_python)
return sys.executable
def _export_session_workflow_in_runtime_env(
workflow_json,
execution_mode="session",
):
"""Export workflow via subprocess in ComfyUI runtime env.
Re-enters the runtime interpreter so ``ComfyUItoPython`` can import
ComfyUI's nodes.py (which requires torch, not available in the test venv).
"""
runtime_python = _get_runtime_python()
env = os.environ.copy()
env["COMFYUI_PATH"] = COMFYUI_PATH
env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")]).rstrip(
os.pathsep
)
tmp_path = tempfile.mktemp(suffix=".py")
try:
result = subprocess.run(
[
runtime_python,
str(Path(__file__).resolve().parents[0] / "runtime" / "run_runtime_validation.py"),
"--internal-export",
"text-to-image",
"--execution-mode",
execution_mode,
"--generated-path",
tmp_path,
],
cwd=ROOT,
env=env,
capture_output=True,
text=True,
)
if result.returncode != 0:
output = (result.stderr or result.stdout or "").strip()
if "Missing runtime dependency" not in output and "ModuleNotFoundError" not in output:
classification = "repo regression"
else:
classification = "environment/setup failure"
raise RuntimeError(
f"Runtime export failed: [{classification}] {output}"
)
# Read generated code from the temp file (stdout is polluted by
# ComfyUI runtime prints such as the sys.path line).
return Path(tmp_path).read_text()
finally:
os.unlink(tmp_path)
def _export_workflow(execution_mode="oneshot"):
"""Export workflow to a string using ComfyUItoPython (unit tests, mock nodes)."""
from comfyui_to_python import ComfyUItoPython
output = StringIO()
ComfyUItoPython(
workflow=json.dumps(TEXT_TO_IMAGE_WORKFLOW),
output_file=output,
node_class_mappings=NODENAMES,
execution_mode=execution_mode,
)
return output.getvalue()
class SessionCodeGenerationTest(unittest.TestCase):
"""Unit tests for session mode code generation."""
def test_oneshot_code_has_bootstrap_helpers_no_session(self):
"""Oneshot mode: bootstrap/cleanup helpers present, no WorkflowSession class."""
generated = _export_workflow("oneshot")
ast.parse(generated)
self.assertIn("bootstrap_comfyui_runtime()", generated)
self.assertIn("cleanup_comfyui_runtime(", generated)
self.assertNotIn("class WorkflowSession", generated)
self.assertNotIn("session.run()", generated)
def test_session_code_has_workflow_session_class_and_session_methods(self):
"""Session mode: WorkflowSession class with run() and close() present."""
generated = _export_workflow("session")
ast.parse(generated)
self.assertIn("class WorkflowSession", generated)
self.assertIn("def run(self)", generated)
self.assertIn("def close(self, unload_models", generated)
self.assertIn("session.run()", generated)
self.assertIn("session.close(", generated)
class SessionModeExecutionTest(unittest.TestCase):
"""E2E tests for session mode script execution."""
@unittest.skipIf(not Path(COMFYUI_PATH).is_dir(), "ComfyUI checkout not available")
def test_oneshot_e2e_text_to_image(self):
"""Oneshot mode: generate and run text-to-image workflow, verify PNG output."""
with tempfile.TemporaryDirectory() as tmpdir:
script_path = Path(tmpdir) / "oneshot.py"
env = os.environ.copy()
env["COMFYUI_PATH"] = COMFYUI_PATH
env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")])
runtime_py = _get_runtime_python()
# Export inside ComfyUI env so node mappings resolve
generated = _export_session_workflow_in_runtime_env(
json.dumps(TEXT_TO_IMAGE_WORKFLOW),
execution_mode="oneshot",
)
script_path.write_text(generated)
# Run
result = subprocess.run(
[runtime_py, str(script_path), "--cpu"],
cwd=ROOT,
env=env,
capture_output=True,
text=True,
timeout=300,
)
self.assertEqual(result.returncode, 0, f"stderr: {result.stderr}")
# Check output
output_dir = Path(COMFYUI_PATH) / "output"
new_outputs = list(output_dir.glob("E2E_text_to_image*.png"))
self.assertTrue(len(new_outputs) > 0, "No PNG output produced")
@unittest.skipIf(not Path(COMFYUI_PATH).is_dir(), "ComfyUI checkout not available")
def test_session_e2e_text_to_image(self):
"""Session mode: generate session-mode script, verify WorkflowSession present and run."""
with tempfile.TemporaryDirectory() as tmpdir:
script_path = Path(tmpdir) / "session.py"
env = os.environ.copy()
env["COMFYUI_PATH"] = COMFYUI_PATH
env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")])
runtime_py = _get_runtime_python()
# Export inside ComfyUI env
generated = _export_session_workflow_in_runtime_env(
json.dumps(TEXT_TO_IMAGE_WORKFLOW),
execution_mode="session",
)
script_path.write_text(generated)
# Verify code structure
self.assertIn("class WorkflowSession", generated)
# Run
result = subprocess.run(
[runtime_py, str(script_path), "--cpu"],
cwd=ROOT,
env=env,
capture_output=True,
text=True,
timeout=300,
)
self.assertEqual(result.returncode, 0, f"stderr: {result.stderr}")
# Check output
output_dir = Path(COMFYUI_PATH) / "output"
new_outputs = list(output_dir.glob("E2E_text_to_image*.png"))
self.assertTrue(len(new_outputs) > 0, "No PNG output produced")
@unittest.skipIf(not Path(COMFYUI_PATH).is_dir(), "ComfyUI checkout not available")
def test_session_e2e_multiple_runs(self):
"""Session mode: generate script that calls session.run() 3x in a row, verify no crash."""
with tempfile.TemporaryDirectory() as tmpdir:
script_path = Path(tmpdir) / "session_multi.py"
env = os.environ.copy()
env["COMFYUI_PATH"] = COMFYUI_PATH
env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")])
runtime_py = _get_runtime_python()
# Export session mode inside ComfyUI env
generated = _export_session_workflow_in_runtime_env(
json.dumps(TEXT_TO_IMAGE_WORKFLOW),
execution_mode="session",
)
# Modify to run 3x
modified = generated.replace(
" session.run()",
" session.run()\n session.run()\n session.run()",
)
script_path.write_text(modified)
# Run
result = subprocess.run(
[runtime_py, str(script_path), "--cpu"],
cwd=ROOT,
env=env,
capture_output=True,
text=True,
timeout=600,
)
self.assertEqual(result.returncode, 0, f"stderr: {result.stderr}")
if __name__ == "__main__":
unittest.main()
tests/test_session_renderer.py
+1 -1
View File
@@ -137,7 +137,7 @@ class SessionRendererTest(unittest.TestCase):
)
generated = output.getvalue()
self.assertIn("def close(self, unload_models: bool = True)", generated)
self.assertIn("def close(self, unload_models: bool | None = None)", generated)
def test_session_mode_oneshot_generates_same_code(self):
workflow = {