fix(hardening): address all PR #157 review findings
Critical:
- _load_module(): wrap exec_module() in try/except, pop from sys.modules
on failure so broken modules aren't cached
- get_comfyui_path(): add _is_comfyui_directory() check to third (CWD walk)
strategy so a non-ComfyUI directory named "ComfyUI" isn't accepted
- get_node_class_mappings(): return cached sys.modules["nodes"] if already
loaded to avoid re-executing nodes.py and losing init_extra_nodes() state
Suggestions:
- import_custom_nodes(): wrap sys.path filtering in try/finally so path is
always restored even if _load_module("server") raises mid-execution
- render.py: auto-discover helpers from generated_helpers.__all__ instead of
a manually maintained list
- pyproject.toml: remove duplicate ruff dev dependency (kept standard
[project.optional-dependencies] form)
Nitpicks:
- Remove dead imports (types, StringIO) from test_import_path_resolution.py
- Trim trailing blank lines in pyproject.toml
- Tone down _load_module docstring: "eliminates ALL" -> "significantly reduces"
Tests (7 new cases):
- TestLoadModuleFailureCleanup: broken module removed on exec failure, retry
starts fresh
- TestLoadModuleTemp: temp modules purged from sys.modules (success + failure)
- TestRepeatedNodeLoad: repeated _load_module("nodes") returns cached copy
- TestGetComfyuiPathThirdStrategy: fake "ComfyUI" dir rejected
- TestImportCustomNodesSysPathRestoration: sys.path restored after crash
This commit is contained in:
@@ -12,9 +12,7 @@ Covers Approach B from docs/specs/harden-import-path-resolution.html:
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from io import StringIO
|
||||
from unittest.mock import patch
|
||||
|
||||
# Real ComfyUI checkout for integration-style unit tests (parent of this repo)
|
||||
@@ -192,7 +190,9 @@ class TestGeneratedScriptIsolation(unittest.TestCase):
|
||||
renderer = WorkflowRenderer()
|
||||
generated = renderer.render(plan)
|
||||
|
||||
self.assertIn("_load_module", generated, "Generated script must embed _load_module()")
|
||||
self.assertIn(
|
||||
"_load_module", generated, "Generated script must embed _load_module()"
|
||||
)
|
||||
self.assertIn(
|
||||
"spec_from_file_location",
|
||||
generated,
|
||||
@@ -283,5 +283,182 @@ class TestIsComfyuiDirectory(unittest.TestCase):
|
||||
self.assertFalse(_is_comfyui_directory("/nonexistent/path/to/comfyui"))
|
||||
|
||||
|
||||
class TestLoadModuleFailureCleanup(unittest.TestCase):
|
||||
"""Tests for _load_module() behavior when exec_module() raises."""
|
||||
|
||||
def tearDown(self):
|
||||
mod_names_to_clean = [
|
||||
name for name in sys.modules if name.startswith("_test_fail_")
|
||||
]
|
||||
for name in mod_names_to_clean:
|
||||
del sys.modules[name]
|
||||
|
||||
def test_broken_module_removed_from_sys_modules_on_failure(self):
|
||||
"""Given a module that raises during exec, it is removed from sys.modules."""
|
||||
from comfyui_to_python.node_runtime import _load_module
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
mod_file = os.path.join(tmpdir, "badmodule.py")
|
||||
with open(mod_file, "w") as f:
|
||||
f.write("raise RuntimeError('mid-execution failure')\n")
|
||||
|
||||
result = _load_module("_test_fail_bad", mod_file)
|
||||
self.assertIsNone(result)
|
||||
self.assertNotIn(
|
||||
"_test_fail_bad",
|
||||
sys.modules,
|
||||
"Failed module must not remain cached in sys.modules",
|
||||
)
|
||||
|
||||
def test_broken_module_retried_cleanly(self):
|
||||
"""Given a failed load, a subsequent load attempt starts fresh (not from cache)."""
|
||||
from comfyui_to_python.node_runtime import _load_module
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
bad_file = os.path.join(tmpdir, "bad.py")
|
||||
good_file = os.path.join(tmpdir, "good.py")
|
||||
|
||||
# First load: bad module that raises
|
||||
with open(bad_file, "w") as f:
|
||||
f.write("raise RuntimeError('fail')\n")
|
||||
result_bad = _load_module("_test_fail_retry", bad_file)
|
||||
self.assertIsNone(result_bad)
|
||||
|
||||
# Overwrite with good module
|
||||
os.remove(bad_file)
|
||||
with open(good_file, "w") as f:
|
||||
f.write("VALUE = 99\n")
|
||||
|
||||
# Reload same name from a different (working) file
|
||||
result_good = _load_module("_test_fail_retry", good_file)
|
||||
self.assertIsNotNone(result_good)
|
||||
self.assertEqual(result_good.VALUE, 99)
|
||||
|
||||
|
||||
class TestLoadModuleTemp(unittest.TestCase):
|
||||
"""Tests for _load_module_temp() removal from sys.modules."""
|
||||
|
||||
def tearDown(self):
|
||||
mod_names_to_clean = [
|
||||
name for name in sys.modules if name.startswith("_test_temp_")
|
||||
]
|
||||
for name in mod_names_to_clean:
|
||||
del sys.modules[name]
|
||||
|
||||
def test_module_removed_from_sys_modules_after_temp_load(self):
|
||||
"""After _load_module_temp(), the module key is absent from sys.modules."""
|
||||
from comfyui_to_python.node_runtime import _load_module_temp
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
mod_file = os.path.join(tmpdir, "tempmod.py")
|
||||
with open(mod_file, "w") as f:
|
||||
f.write("VALUE = 123\n")
|
||||
|
||||
mod = _load_module_temp("_test_temp_mod", mod_file)
|
||||
self.assertIsNotNone(mod)
|
||||
self.assertEqual(mod.VALUE, 123)
|
||||
self.assertNotIn(
|
||||
"_test_temp_mod",
|
||||
sys.modules,
|
||||
"Temp module must be removed from sys.modules after load",
|
||||
)
|
||||
|
||||
def test_module_removed_on_temp_load_failure(self):
|
||||
"""After _load_module_temp() fails, the module key is absent from sys.modules."""
|
||||
from comfyui_to_python.node_runtime import _load_module_temp
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
mod_file = os.path.join(tmpdir, "badtemp.py")
|
||||
with open(mod_file, "w") as f:
|
||||
f.write("raise RuntimeError('fail')\n")
|
||||
|
||||
result = _load_module_temp("_test_temp_bad", mod_file)
|
||||
self.assertIsNone(result)
|
||||
self.assertNotIn(
|
||||
"_test_temp_bad",
|
||||
sys.modules,
|
||||
"Failed temp module must be removed from sys.modules",
|
||||
)
|
||||
|
||||
|
||||
class TestRepeatedNodeLoad(unittest.TestCase):
|
||||
"""Tests that repeated _load_module('nodes', ...) calls respect the cache."""
|
||||
|
||||
def tearDown(self):
|
||||
if "_test_nodes" in sys.modules:
|
||||
del sys.modules["_test_nodes"]
|
||||
|
||||
def test_repeated_load_preserves_cached_attributes(self):
|
||||
"""Second _load_module() for same name returns cached module, not a fresh copy."""
|
||||
from comfyui_to_python.node_runtime import _load_module
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
mod_file = os.path.join(tmpdir, "mynodes.py")
|
||||
with open(mod_file, "w") as f:
|
||||
f.write("NODE_CLASS_MAPPINGS = {'A': 1}\n")
|
||||
|
||||
first = _load_module("_test_nodes", mod_file)
|
||||
self.assertIsNotNone(first)
|
||||
|
||||
# Mutate the cached module to simulate init_extra_nodes()
|
||||
first.NODE_CLASS_MAPPINGS["B"] = 2
|
||||
|
||||
# Second load with same name must return cached copy, preserving mutation
|
||||
second = _load_module("_test_nodes", mod_file)
|
||||
self.assertIs(first, second)
|
||||
self.assertEqual(second.NODE_CLASS_MAPPINGS.get("B"), 2)
|
||||
|
||||
|
||||
class TestGetComfyuiPathThirdStrategy(unittest.TestCase):
|
||||
"""Tests for get_comfyui_path() third-strategy verification."""
|
||||
|
||||
def test_fallback_rejects_non_comfyui_directory(self):
|
||||
"""Given a directory named ComfyUI without nodes.py, third strategy rejects it."""
|
||||
from comfyui_to_python.node_runtime import get_comfyui_path
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# Create a fake "ComfyUI" directory (no nodes.py)
|
||||
fake_dir = os.path.join(tmpdir, "ComfyUI")
|
||||
os.makedirs(fake_dir)
|
||||
|
||||
original_cwd = os.getcwd()
|
||||
try:
|
||||
os.chdir(tmpdir)
|
||||
with patch.dict("os.environ", {}, clear=False):
|
||||
# Strip COMFYUI_PATH so fallback strategies activate
|
||||
env_copy = dict(os.environ)
|
||||
env_copy.pop("COMFYUI_PATH", None)
|
||||
with patch.dict("os.environ", env_copy, clear=True):
|
||||
result = get_comfyui_path()
|
||||
finally:
|
||||
os.chdir(original_cwd)
|
||||
|
||||
# Must NOT return the fake directory
|
||||
self.assertNotEqual(result, fake_dir)
|
||||
|
||||
|
||||
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."""
|
||||
from comfyui_to_python.node_runtime import import_custom_nodes
|
||||
|
||||
original_path = sys.path[:]
|
||||
|
||||
with patch(
|
||||
"comfyui_to_python.node_runtime._load_module",
|
||||
side_effect=RuntimeError("simulated crash during server load"),
|
||||
):
|
||||
try:
|
||||
import_custom_nodes()
|
||||
except RuntimeError:
|
||||
pass # Expected crash
|
||||
|
||||
# sys.path should still contain all original entries (no dangling state)
|
||||
for entry in original_path:
|
||||
self.assertIn(entry, sys.path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user