fix(hardening): address PR #157 review findings
Critical fixes: - node_runtime.py: change outer handler in _load_module() from except Exception to except BaseException with sys.modules.pop() cleanup, preventing non-Exception BaseExceptions (SystemExit, KeyboardInterrupt) from violating the no-raise contract - app.py: remove incorrect reassignment of base_node_class_mappings after custom node import; keep it as pre-custom-node baseline so WorkflowPlanner correctly treats extras/custom nodes as NODE_CLASS_MAPPINGS dict lookups instead of direct imports - Regenerate stale test fixtures (text-to-image.py, upscale-model-loader.py) from current source: updated _load_module() with exec failure cleanup pattern, 3-marker structural verification, and try/finally sys.path restoration in import_custom_nodes() Suggestions addressed: - node_runtime.py: _is_comfyui_directory() now checks 3 structural markers (nodes.py, main.py, comfy/) instead of just nodes.py - generator/render.py: added None guard for getattr on generated_helpers.__all__ items with log warning for missing helpers - test_import_path_resolution.py: fixed test_sys_path_restored_on_crash to target the server module specifically after sys.path manipulation, so the try/finally restoration path is actually exercised Nits addressed: - .gitignore: added missing trailing newline New test coverage (tests/test_app_base_mappings.py): - base_node_class_mappings unchanged after custom node import - deep copy independence of base_node_class_mappings - render output contains hardened _load_module with except BaseException pattern
This commit is contained in:
Peyton DeNiro
+1
-1
@@ -149,4 +149,4 @@ cython_debug/
|
||||
docs/
|
||||
.agents/
|
||||
.codex/
|
||||
AGENTS.md
|
||||
AGENTS.md
|
||||
|
||||
@@ -65,7 +65,9 @@ class ExportApplication:
|
||||
if nodes_mod is not None:
|
||||
fresh_mappings = getattr(nodes_mod, "NODE_CLASS_MAPPINGS", {})
|
||||
self.node_class_mappings = fresh_mappings
|
||||
self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings)
|
||||
# Leave self.base_node_class_mappings unchanged — it represents the
|
||||
# pre-custom-node baseline used by WorkflowPlanner to decide whether
|
||||
# a node gets a direct import or a NODE_CLASS_MAPPINGS dict lookup.
|
||||
|
||||
load_order = LoadOrderDeterminer(
|
||||
data, self.node_class_mappings
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
import inspect
|
||||
import logging
|
||||
from pprint import pformat
|
||||
from typing import Any
|
||||
|
||||
import black
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
from ..node_runtime import import_custom_nodes
|
||||
from . import generated_helpers
|
||||
from .model import GenerationPlan
|
||||
@@ -25,7 +28,12 @@ class WorkflowRenderer:
|
||||
# are picked up without manually updating this list.
|
||||
func_strings = []
|
||||
for name in generated_helpers.__all__:
|
||||
func = getattr(generated_helpers, name)
|
||||
func = getattr(generated_helpers, name, None)
|
||||
if func is None:
|
||||
log.warning(
|
||||
"Helper '%s' missing from generated_helpers — skipping", name
|
||||
)
|
||||
continue
|
||||
func_strings.append(f"\n{inspect.getsource(func)}")
|
||||
|
||||
static_imports = [
|
||||
|
||||
@@ -9,10 +9,18 @@ log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _is_comfyui_directory(path: str) -> bool:
|
||||
"""Verify a directory has ComfyUI structural markers (nodes.py)."""
|
||||
"""Verify a directory has ComfyUI structural markers.
|
||||
|
||||
Checks for nodes.py, main.py, and the comfy/ subdirectory to raise
|
||||
the bar against spoofing via a directory with only a single marker file.
|
||||
"""
|
||||
if not os.path.isdir(path):
|
||||
return False
|
||||
return os.path.isfile(os.path.join(path, "nodes.py"))
|
||||
return (
|
||||
os.path.isfile(os.path.join(path, "nodes.py"))
|
||||
and os.path.isfile(os.path.join(path, "main.py"))
|
||||
and os.path.isdir(os.path.join(path, "comfy"))
|
||||
)
|
||||
|
||||
|
||||
def _load_module(module_name: str, filepath: str) -> Any:
|
||||
@@ -43,8 +51,9 @@ def _load_module(module_name: str, filepath: str) -> Any:
|
||||
sys.modules.pop(module_name, None)
|
||||
raise
|
||||
return mod
|
||||
except Exception as e:
|
||||
except BaseException as e:
|
||||
log.debug("Failed to load %s from %s: %s", module_name, filepath, e)
|
||||
sys.modules.pop(module_name, None) # Also clean up on edge-case exceptions
|
||||
return None
|
||||
|
||||
|
||||
|
||||
@@ -10,47 +10,6 @@ from typing import Sequence, Mapping, Any, Union
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _load_module(module_name: str, filepath: str) -> Any:
|
||||
"""Load a Python module from an explicit file path, bypassing sys.path.
|
||||
|
||||
This eliminates ALL bare import shadowing attacks by loading modules
|
||||
from verified file paths instead of relying on sys.path resolution.
|
||||
"""
|
||||
if not os.path.isfile(filepath):
|
||||
log.debug("Module file not found: %s (%s)", module_name, filepath)
|
||||
return None
|
||||
try:
|
||||
spec = importlib.util.spec_from_file_location(module_name, filepath)
|
||||
if spec is None or spec.loader is None:
|
||||
log.debug("Could not create spec for %s at %s", module_name, filepath)
|
||||
return None
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = mod
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
except Exception as e:
|
||||
log.debug("Failed to load %s from %s: %s", module_name, filepath, e)
|
||||
return None
|
||||
|
||||
|
||||
def _load_module_temp(module_name: str, filepath: str) -> Any:
|
||||
"""Load a module via _load_module() then remove it from sys.modules.
|
||||
|
||||
Used during bootstrap for modules that ComfyUI's import chain also loads
|
||||
normally — prevents the cached copy from conflicting with later imports.
|
||||
"""
|
||||
mod = _load_module(module_name, filepath)
|
||||
sys.modules.pop(module_name, None)
|
||||
return mod
|
||||
|
||||
|
||||
def _is_comfyui_directory(path: str) -> bool:
|
||||
"""Verify a directory has ComfyUI structural markers (nodes.py)."""
|
||||
if not os.path.isdir(path):
|
||||
return False
|
||||
return os.path.isfile(os.path.join(path, "nodes.py"))
|
||||
|
||||
|
||||
def _find_from_extension_location() -> str | None:
|
||||
"""Walk up from this file's location to find ComfyUI root."""
|
||||
ext_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
@@ -66,46 +25,64 @@ def _find_from_extension_location() -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
|
||||
"""Return a sequence or mapping result item by index."""
|
||||
def _is_comfyui_directory(path: str) -> bool:
|
||||
"""Verify a directory has ComfyUI structural markers.
|
||||
|
||||
Checks for nodes.py, main.py, and the comfy/ subdirectory to raise
|
||||
the bar against spoofing via a directory with only a single marker file.
|
||||
"""
|
||||
if not os.path.isdir(path):
|
||||
return False
|
||||
return (
|
||||
os.path.isfile(os.path.join(path, "nodes.py"))
|
||||
and os.path.isfile(os.path.join(path, "main.py"))
|
||||
and os.path.isdir(os.path.join(path, "comfy"))
|
||||
)
|
||||
|
||||
|
||||
def _load_module(module_name: str, filepath: str) -> Any:
|
||||
"""Load a Python module from an explicit file path, bypassing sys.path.
|
||||
|
||||
Significantly reduces bare import shadowing risk by loading modules
|
||||
from verified file paths instead of relying on sys.path resolution.
|
||||
If exec_module() raises, the partially-loaded module is removed from
|
||||
sys.modules so subsequent calls start fresh.
|
||||
"""
|
||||
# Return cached module if already loaded — prevents re-execution
|
||||
# that would reset state (e.g. NODE_CLASS_MAPPINGS after init_extra_nodes).
|
||||
if module_name in sys.modules:
|
||||
return sys.modules[module_name]
|
||||
if not os.path.isfile(filepath):
|
||||
log.debug("Module file not found: %s (%s)", module_name, filepath)
|
||||
return None
|
||||
try:
|
||||
return obj[index]
|
||||
except KeyError:
|
||||
return obj["result"][index]
|
||||
spec = importlib.util.spec_from_file_location(module_name, filepath)
|
||||
if spec is None or spec.loader is None:
|
||||
log.debug("Could not create spec for %s at %s", module_name, filepath)
|
||||
return None
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = mod
|
||||
try:
|
||||
spec.loader.exec_module(mod)
|
||||
except BaseException:
|
||||
sys.modules.pop(module_name, None)
|
||||
raise
|
||||
return mod
|
||||
except BaseException as e:
|
||||
log.debug("Failed to load %s from %s: %s", module_name, filepath, e)
|
||||
sys.modules.pop(module_name, None) # Also clean up on edge-case exceptions
|
||||
return None
|
||||
|
||||
|
||||
def get_comfyui_path() -> str | None:
|
||||
"""Resolve ComfyUI path via prioritized multi-strategy fallback.
|
||||
def _load_module_temp(module_name: str, filepath: str) -> Any:
|
||||
"""Load a module via _load_module() then remove it from sys.modules.
|
||||
|
||||
Strategy order:
|
||||
1. COMFYUI_PATH env var (verified with _is_comfyui_directory)
|
||||
2. Relative walk from extension location (realpath + verified)
|
||||
3. CWD walk (legacy fallback, depth-limited)
|
||||
Used during bootstrap for modules that ComfyUI's import chain also loads
|
||||
normally — prevents the cached copy from conflicting with later imports.
|
||||
"""
|
||||
p = os.environ.get("COMFYUI_PATH")
|
||||
if p and _is_comfyui_directory(p):
|
||||
return p
|
||||
p = _find_from_extension_location()
|
||||
if p:
|
||||
return p
|
||||
return find_path("ComfyUI", max_depth=20)
|
||||
|
||||
|
||||
def find_path(name: str, max_depth: int = 20) -> str | None:
|
||||
"""Iteratively walk up from CWD to find a directory by name.
|
||||
|
||||
Depth-limited to prevent slow startup on deep trees.
|
||||
Each candidate verified by _is_comfyui_directory() at the caller.
|
||||
"""
|
||||
candidate = os.getcwd()
|
||||
for _ in range(max_depth):
|
||||
parent = os.path.dirname(candidate)
|
||||
if parent == candidate:
|
||||
break
|
||||
candidate = parent
|
||||
if os.path.basename(candidate) == name:
|
||||
return candidate
|
||||
return None
|
||||
mod = _load_module(module_name, filepath)
|
||||
sys.modules.pop(module_name, None)
|
||||
return mod
|
||||
|
||||
|
||||
def add_comfyui_directory_to_sys_path() -> None:
|
||||
@@ -247,6 +224,51 @@ def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None:
|
||||
gc.collect()
|
||||
|
||||
|
||||
def find_path(name: str, max_depth: int = 20) -> str | None:
|
||||
"""Iteratively walk up from CWD to find a directory by name.
|
||||
|
||||
Depth-limited to prevent slow startup on deep trees.
|
||||
Each candidate verified by _is_comfyui_directory() at the caller.
|
||||
"""
|
||||
candidate = os.getcwd()
|
||||
for _ in range(max_depth):
|
||||
parent = os.path.dirname(candidate)
|
||||
if parent == candidate:
|
||||
break
|
||||
candidate = parent
|
||||
if os.path.basename(candidate) == name:
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def get_comfyui_path() -> str | None:
|
||||
"""Resolve ComfyUI path via prioritized multi-strategy fallback.
|
||||
|
||||
Strategy order:
|
||||
1. COMFYUI_PATH env var (verified with _is_comfyui_directory)
|
||||
2. Relative walk from extension location (realpath + verified)
|
||||
3. CWD walk (legacy fallback, depth-limited)
|
||||
"""
|
||||
p = os.environ.get("COMFYUI_PATH")
|
||||
if p and _is_comfyui_directory(p):
|
||||
return p
|
||||
p = _find_from_extension_location()
|
||||
if p:
|
||||
return p
|
||||
p = find_path("ComfyUI", max_depth=20)
|
||||
if p and _is_comfyui_directory(p):
|
||||
return p
|
||||
return None
|
||||
|
||||
|
||||
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
|
||||
"""Return a sequence or mapping result item by index."""
|
||||
try:
|
||||
return obj[index]
|
||||
except KeyError:
|
||||
return obj["result"][index]
|
||||
|
||||
|
||||
# Workflow data
|
||||
def build_workflow() -> dict[str, Any]:
|
||||
return {
|
||||
|
||||
@@ -10,47 +10,6 @@ from typing import Sequence, Mapping, Any, Union
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _load_module(module_name: str, filepath: str) -> Any:
|
||||
"""Load a Python module from an explicit file path, bypassing sys.path.
|
||||
|
||||
This eliminates ALL bare import shadowing attacks by loading modules
|
||||
from verified file paths instead of relying on sys.path resolution.
|
||||
"""
|
||||
if not os.path.isfile(filepath):
|
||||
log.debug("Module file not found: %s (%s)", module_name, filepath)
|
||||
return None
|
||||
try:
|
||||
spec = importlib.util.spec_from_file_location(module_name, filepath)
|
||||
if spec is None or spec.loader is None:
|
||||
log.debug("Could not create spec for %s at %s", module_name, filepath)
|
||||
return None
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = mod
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
except Exception as e:
|
||||
log.debug("Failed to load %s from %s: %s", module_name, filepath, e)
|
||||
return None
|
||||
|
||||
|
||||
def _load_module_temp(module_name: str, filepath: str) -> Any:
|
||||
"""Load a module via _load_module() then remove it from sys.modules.
|
||||
|
||||
Used during bootstrap for modules that ComfyUI's import chain also loads
|
||||
normally — prevents the cached copy from conflicting with later imports.
|
||||
"""
|
||||
mod = _load_module(module_name, filepath)
|
||||
sys.modules.pop(module_name, None)
|
||||
return mod
|
||||
|
||||
|
||||
def _is_comfyui_directory(path: str) -> bool:
|
||||
"""Verify a directory has ComfyUI structural markers (nodes.py)."""
|
||||
if not os.path.isdir(path):
|
||||
return False
|
||||
return os.path.isfile(os.path.join(path, "nodes.py"))
|
||||
|
||||
|
||||
def _find_from_extension_location() -> str | None:
|
||||
"""Walk up from this file's location to find ComfyUI root."""
|
||||
ext_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
@@ -66,46 +25,64 @@ def _find_from_extension_location() -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
|
||||
"""Return a sequence or mapping result item by index."""
|
||||
def _is_comfyui_directory(path: str) -> bool:
|
||||
"""Verify a directory has ComfyUI structural markers.
|
||||
|
||||
Checks for nodes.py, main.py, and the comfy/ subdirectory to raise
|
||||
the bar against spoofing via a directory with only a single marker file.
|
||||
"""
|
||||
if not os.path.isdir(path):
|
||||
return False
|
||||
return (
|
||||
os.path.isfile(os.path.join(path, "nodes.py"))
|
||||
and os.path.isfile(os.path.join(path, "main.py"))
|
||||
and os.path.isdir(os.path.join(path, "comfy"))
|
||||
)
|
||||
|
||||
|
||||
def _load_module(module_name: str, filepath: str) -> Any:
|
||||
"""Load a Python module from an explicit file path, bypassing sys.path.
|
||||
|
||||
Significantly reduces bare import shadowing risk by loading modules
|
||||
from verified file paths instead of relying on sys.path resolution.
|
||||
If exec_module() raises, the partially-loaded module is removed from
|
||||
sys.modules so subsequent calls start fresh.
|
||||
"""
|
||||
# Return cached module if already loaded — prevents re-execution
|
||||
# that would reset state (e.g. NODE_CLASS_MAPPINGS after init_extra_nodes).
|
||||
if module_name in sys.modules:
|
||||
return sys.modules[module_name]
|
||||
if not os.path.isfile(filepath):
|
||||
log.debug("Module file not found: %s (%s)", module_name, filepath)
|
||||
return None
|
||||
try:
|
||||
return obj[index]
|
||||
except KeyError:
|
||||
return obj["result"][index]
|
||||
spec = importlib.util.spec_from_file_location(module_name, filepath)
|
||||
if spec is None or spec.loader is None:
|
||||
log.debug("Could not create spec for %s at %s", module_name, filepath)
|
||||
return None
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = mod
|
||||
try:
|
||||
spec.loader.exec_module(mod)
|
||||
except BaseException:
|
||||
sys.modules.pop(module_name, None)
|
||||
raise
|
||||
return mod
|
||||
except BaseException as e:
|
||||
log.debug("Failed to load %s from %s: %s", module_name, filepath, e)
|
||||
sys.modules.pop(module_name, None) # Also clean up on edge-case exceptions
|
||||
return None
|
||||
|
||||
|
||||
def get_comfyui_path() -> str | None:
|
||||
"""Resolve ComfyUI path via prioritized multi-strategy fallback.
|
||||
def _load_module_temp(module_name: str, filepath: str) -> Any:
|
||||
"""Load a module via _load_module() then remove it from sys.modules.
|
||||
|
||||
Strategy order:
|
||||
1. COMFYUI_PATH env var (verified with _is_comfyui_directory)
|
||||
2. Relative walk from extension location (realpath + verified)
|
||||
3. CWD walk (legacy fallback, depth-limited)
|
||||
Used during bootstrap for modules that ComfyUI's import chain also loads
|
||||
normally — prevents the cached copy from conflicting with later imports.
|
||||
"""
|
||||
p = os.environ.get("COMFYUI_PATH")
|
||||
if p and _is_comfyui_directory(p):
|
||||
return p
|
||||
p = _find_from_extension_location()
|
||||
if p:
|
||||
return p
|
||||
return find_path("ComfyUI", max_depth=20)
|
||||
|
||||
|
||||
def find_path(name: str, max_depth: int = 20) -> str | None:
|
||||
"""Iteratively walk up from CWD to find a directory by name.
|
||||
|
||||
Depth-limited to prevent slow startup on deep trees.
|
||||
Each candidate verified by _is_comfyui_directory() at the caller.
|
||||
"""
|
||||
candidate = os.getcwd()
|
||||
for _ in range(max_depth):
|
||||
parent = os.path.dirname(candidate)
|
||||
if parent == candidate:
|
||||
break
|
||||
candidate = parent
|
||||
if os.path.basename(candidate) == name:
|
||||
return candidate
|
||||
return None
|
||||
mod = _load_module(module_name, filepath)
|
||||
sys.modules.pop(module_name, None)
|
||||
return mod
|
||||
|
||||
|
||||
def add_comfyui_directory_to_sys_path() -> None:
|
||||
@@ -247,6 +224,54 @@ def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None:
|
||||
gc.collect()
|
||||
|
||||
|
||||
def find_path(name: str, max_depth: int = 20) -> str | None:
|
||||
"""Iteratively walk up from CWD to find a directory by name.
|
||||
|
||||
Depth-limited to prevent slow startup on deep trees.
|
||||
Each candidate verified by _is_comfyui_directory() at the caller.
|
||||
"""
|
||||
candidate = os.getcwd()
|
||||
for _ in range(max_depth):
|
||||
parent = os.path.dirname(candidate)
|
||||
if parent == candidate:
|
||||
break
|
||||
candidate = parent
|
||||
if os.path.basename(candidate) == name:
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def get_comfyui_path() -> str | None:
|
||||
"""Resolve ComfyUI path via prioritized multi-strategy fallback.
|
||||
|
||||
Strategy order:
|
||||
1. COMFYUI_PATH env var (verified with _is_comfyui_directory)
|
||||
2. Relative walk from extension location (realpath + verified)
|
||||
3. CWD walk (legacy fallback, depth-limited)
|
||||
"""
|
||||
p = os.environ.get("COMFYUI_PATH")
|
||||
if p and _is_comfyui_directory(p):
|
||||
return p
|
||||
p = _find_from_extension_location()
|
||||
if p:
|
||||
return p
|
||||
p = find_path("ComfyUI", max_depth=20)
|
||||
if p and _is_comfyui_directory(p):
|
||||
return p
|
||||
return None
|
||||
|
||||
|
||||
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
|
||||
"""Return a sequence or mapping result item by index."""
|
||||
try:
|
||||
return obj[index]
|
||||
except KeyError:
|
||||
return obj["result"][index]
|
||||
|
||||
|
||||
import_custom_nodes()
|
||||
|
||||
|
||||
def import_custom_nodes() -> None:
|
||||
"""Initialize ComfyUI custom nodes in the exporter runtime.
|
||||
|
||||
@@ -277,10 +302,12 @@ def import_custom_nodes() -> None:
|
||||
original_sys_path = list(sys.path)
|
||||
sys.path[:] = [p for p in sys.path if p != comfy_subdir]
|
||||
|
||||
server_mod = _load_module("server", os.path.join(comfyui_path, "server.py"))
|
||||
|
||||
# Restore sys.path so nodes and other modules that need comfy/ still work
|
||||
sys.path[:] = original_sys_path
|
||||
try:
|
||||
server_mod = _load_module("server", os.path.join(comfyui_path, "server.py"))
|
||||
finally:
|
||||
# Restore sys.path so nodes and other modules that need comfy/ still work.
|
||||
# Guaranteed even if _load_module raises mid-execution.
|
||||
sys.path[:] = original_sys_path
|
||||
|
||||
if execution_mod is None or server_mod is None:
|
||||
log.debug(
|
||||
|
||||
@@ -0,0 +1,165 @@
|
||||
"""Tests for ExportApplication base_node_class_mappings behavior.
|
||||
|
||||
Verifies that after custom node import, base_node_class_mappings retains the
|
||||
pre-custom-node baseline so WorkflowPlanner correctly distinguishes between
|
||||
built-in nodes (direct import) and custom/extras nodes (NODE_CLASS_MAPPINGS lookup).
|
||||
"""
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
from io import StringIO
|
||||
|
||||
|
||||
def _make_workflow():
|
||||
"""Minimal workflow JSON string for testing."""
|
||||
return '{"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "test.safetensors"}}}'
|
||||
|
||||
|
||||
class TestBaseNodeClassMappings(unittest.TestCase):
|
||||
"""Tests for base_node_class_mappings stability across custom node reloads."""
|
||||
|
||||
def tearDown(self):
|
||||
# Clean up any sys.modules pollution to avoid cross-test contamination
|
||||
if "nodes" in sys.modules:
|
||||
del sys.modules["nodes"]
|
||||
|
||||
def test_base_mappings_unchanged_after_custom_node_import(self):
|
||||
"""Given custom nodes are loaded, base_node_class_mappings stays at original baseline."""
|
||||
from comfyui_to_python.app import ExportApplication
|
||||
|
||||
class FakeNode:
|
||||
CATEGORY = "loaders"
|
||||
FUNCTION = "execute"
|
||||
|
||||
@staticmethod
|
||||
def INPUT_TYPES():
|
||||
return {"required": {"ckpt_name": ("STRING",)}}
|
||||
|
||||
def execute(self, ckpt_name: str):
|
||||
return ({},)
|
||||
|
||||
class FakeCustomNode:
|
||||
CATEGORY = "custom"
|
||||
FUNCTION = "execute"
|
||||
|
||||
@staticmethod
|
||||
def INPUT_TYPES():
|
||||
return {"required": {"ckpt_name": ("STRING",)}}
|
||||
|
||||
def execute(self, ckpt_name: str):
|
||||
return ({},)
|
||||
|
||||
initial_mappings = {"CheckpointLoaderSimple": FakeNode}
|
||||
after_import_mappings = dict(initial_mappings)
|
||||
# Simulate custom nodes adding a new class
|
||||
after_import_mappings["CustomNode"] = FakeCustomNode
|
||||
|
||||
call_count = [0]
|
||||
|
||||
def mock_node_loader():
|
||||
return initial_mappings
|
||||
|
||||
def mock_custom_node_importer():
|
||||
call_count[0] += 1
|
||||
# Simulate populating sys.modules["nodes"].NODE_CLASS_MAPPINGS
|
||||
nodes_mod_type = type("NodesModule", (), {})()
|
||||
nodes_mod_type.NODE_CLASS_MAPPINGS = after_import_mappings
|
||||
sys.modules["nodes"] = nodes_mod_type
|
||||
|
||||
output = StringIO()
|
||||
app = ExportApplication(
|
||||
workflow=_make_workflow(),
|
||||
output_file=output,
|
||||
node_mapping_loader=mock_node_loader,
|
||||
custom_node_importer=mock_custom_node_importer,
|
||||
needs_init_custom_nodes=True,
|
||||
)
|
||||
|
||||
# base_node_class_mappings should be a deep copy of initial (pre-custom-node) state
|
||||
self.assertEqual(app.base_node_class_mappings, initial_mappings)
|
||||
app.execute()
|
||||
|
||||
# After execute(), base_node_class_mappings must still be the original baseline
|
||||
# (not overwritten with after_import_mappings which includes custom nodes)
|
||||
self.assertEqual(
|
||||
app.base_node_class_mappings,
|
||||
initial_mappings,
|
||||
"base_node_class_mappings should remain at pre-custom-node baseline",
|
||||
)
|
||||
|
||||
def test_base_mappings_is_deep_copy_not_shared_reference(self):
|
||||
"""Given node mappings change in-place, base_node_class_mappings is independent."""
|
||||
from comfyui_to_python.app import ExportApplication
|
||||
|
||||
class FakeNode2:
|
||||
CATEGORY = "loaders"
|
||||
FUNCTION = "execute"
|
||||
|
||||
@staticmethod
|
||||
def INPUT_TYPES():
|
||||
return {"required": {"ckpt_name": ("STRING",)}}
|
||||
|
||||
def execute(self, ckpt_name: str):
|
||||
return ({},)
|
||||
|
||||
initial_mappings = {"CheckpointLoaderSimple": FakeNode2}
|
||||
|
||||
def mock_node_loader():
|
||||
return dict(initial_mappings)
|
||||
|
||||
output = StringIO()
|
||||
app = ExportApplication(
|
||||
workflow=_make_workflow(),
|
||||
output_file=output,
|
||||
node_mapping_loader=mock_node_loader,
|
||||
custom_node_importer=lambda: None,
|
||||
)
|
||||
|
||||
# Mutate the live mappings
|
||||
app.node_class_mappings["NewKey"] = "value"
|
||||
|
||||
# base_node_class_mappings must not reflect the mutation (deep copy)
|
||||
self.assertNotIn(
|
||||
"NewKey",
|
||||
app.base_node_class_mappings,
|
||||
"base_node_class_mappings must be a deep copy, not shared reference",
|
||||
)
|
||||
|
||||
|
||||
class TestRenderHardenedImports(unittest.TestCase):
|
||||
"""Integration test: render output contains hardened _load_module pattern."""
|
||||
|
||||
def test_generated_output_contains_exec_failure_cleanup(self):
|
||||
"""Given a GenerationPlan, rendered output includes try/except BaseException around exec_module()."""
|
||||
from comfyui_to_python.generator.render import WorkflowRenderer
|
||||
from comfyui_to_python.generator.model import GenerationPlan
|
||||
|
||||
plan = GenerationPlan(
|
||||
workflow_data={"1": {"class_type": "CheckpointLoaderSimple"}},
|
||||
metadata_workflow_data=None,
|
||||
custom_nodes=False,
|
||||
import_statements={},
|
||||
special_functions_code=[],
|
||||
loop_code=["result = some_node()"],
|
||||
queue_size=1,
|
||||
)
|
||||
renderer = WorkflowRenderer()
|
||||
generated = renderer.render(plan)
|
||||
|
||||
# The _load_module in rendered output must have the exec failure cleanup pattern
|
||||
self.assertIn("try:", generated)
|
||||
self.assertIn("spec.loader.exec_module(mod)", generated)
|
||||
self.assertIn(
|
||||
"except BaseException:",
|
||||
generated,
|
||||
"Must catch BaseException for edge-case safety",
|
||||
)
|
||||
self.assertIn(
|
||||
"sys.modules.pop(module_name, None)",
|
||||
generated,
|
||||
"Must clean up sys.modules on exec failure",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -441,20 +441,40 @@ class TestImportCustomNodesSysPathRestoration(unittest.TestCase):
|
||||
"""Tests for sys.path restoration in import_custom_nodes()."""
|
||||
|
||||
def test_sys_path_restored_on_crash(self):
|
||||
"""Given a crash during import_custom_nodes(), sys.path is still restored."""
|
||||
"""Given a crash during import_custom_nodes(), sys.path is still restored.
|
||||
|
||||
The crash fires on the "server" module load (after comfy/ filtering)
|
||||
so that the try/finally sys.path restoration path is actually exercised.
|
||||
"""
|
||||
from comfyui_to_python.node_runtime import import_custom_nodes
|
||||
|
||||
original_path = sys.path[:]
|
||||
call_order = []
|
||||
|
||||
def fake_load(name, path):
|
||||
call_order.append(name)
|
||||
if name == "server":
|
||||
raise RuntimeError("simulated crash during server load")
|
||||
return None # Let execution/nodes return None (falls through gracefully)
|
||||
|
||||
with patch(
|
||||
"comfyui_to_python.node_runtime._load_module",
|
||||
side_effect=RuntimeError("simulated crash during server load"),
|
||||
side_effect=fake_load,
|
||||
), patch(
|
||||
"comfyui_to_python.node_runtime.get_comfyui_path",
|
||||
return_value=_REAL_COMFYUI if os.path.isdir(_REAL_COMFYUI) else None,
|
||||
):
|
||||
try:
|
||||
import_custom_nodes()
|
||||
except RuntimeError:
|
||||
pass # Expected crash
|
||||
|
||||
# When ComfyUI path is found, "execution" and "nodes" should be called
|
||||
# before reaching the server module load inside try/finally.
|
||||
self.assertIn("execution", call_order, "Expected 'execution' to be loaded")
|
||||
self.assertIn("nodes", call_order, "Expected 'nodes' to be loaded")
|
||||
self.assertIn("server", call_order, "Expected 'server' to trigger the crash")
|
||||
|
||||
# sys.path should still contain all original entries (no dangling state)
|
||||
for entry in original_path:
|
||||
self.assertIn(entry, sys.path)
|
||||
|
||||
Reference in New Issue
Block a user