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

harden import path resolution: full importlib isolation (Approach B)

Browse Source 
Replace all bare ComfyUI internal imports with centralized _load_module()
using importlib.util.spec_from_file_location(). This eliminates sys.path
shadowing attacks, removes the remove/re-insert gap window, and makes
import resolution deterministic regardless of working directory, install
layout, or prior sys.path state.

node_runtime.py — core rewrite:
- Add _load_module(): centralized importlib isolation layer that loads
  modules from explicit file paths verified against a ComfyUI root
  confirmed by _is_comfyui_directory() (checks for nodes.py marker)
- Add _is_comfyui_directory(): structural verification of candidate paths
- Add _find_from_extension_location(): walk up from __file__ via realpath
  to find ComfyUI root (handles symlinked installs)
- Rewrite get_comfyui_path(): multi-strategy fallback chain:
  (1) COMFYUI_PATH env var with verification, (2) relative extension walk,
  (3) CWD walk as depth-limited last resort (max 20 levels)
- Replace bare 'from nodes import NODE_CLASS_MAPPINGS' with _load_module()
- Replace bare 'import execution/server' with _load_module()
- Replace bare 'from main import load_extra_path_config' with _load_module()
  (alias "comfy_main" to prevent sys.modules cache collision)
- Replace bare 'from utils.extra_config import ...' with _load_module()
- Replace bare 'import comfy.options/cli_args' with _load_module()
- Replace bare 'import cuda_malloc' with _load_module()
- Rewrite add_comfyui_directory_to_sys_path(): idempotent insert-once,
  no remove/re-insert gap window
- Rewrite import_custom_nodes(): _load_module() for execution/nodes/server,
  eliminated sys.path.remove() that created a gap window for lazy imports
- Rewrite get_node_class_mappings(): _load_module() for nodes.py
- Replace all print() statements with logging.DEBUG (no path leaks to stdout)
- Add None guards on args attributes after cli_args module load

generator/render.py — generated script isolation:
- Embed _load_module(), _is_comfyui_directory(), _find_from_extension_location()
  in generated scripts via inspect.getsource()
- Add importlib.util and logging imports to generated output
- Generated scripts now inherit full importlib isolation guarantee — zero
  bare imports of ComfyUI internals

generator/generated_helpers.py:
- Export _load_module, _is_comfyui_directory, _find_from_extension_location
  for use by render.py

tests/test_import_path_resolution.py — new (14 test cases):
- TestIsComfyuiDirectory: structural verification against real ComfyUI checkout,
  rejection of empty/nonexistent directories
- TestLoadModule: module loading from explicit path, graceful None for missing
- TestShadowingResistance: malicious main.py and nodes.py at sys.path[0] are
  ignored; _load_module() always loads from verified file path
- TestSysPathIdempotence: insert-once behavior, no remove/re-insert in
  import_custom_nodes(), no print() statements in module source
- TestGeneratedScriptIsolation: generated output contains _load_module() and
  spec_from_file_location; zero bare ComfyUI imports
- TestGetComfyuiPath: env var strategy with verification, fallback chain

tests/test_upscale_model_loader_export.py — updated assertions:
- Replace "import comfy.options" / "import cuda_malloc" checks with
  _load_module() string pattern matching ("comfy.options", "cuda_malloc")
- Update ordering assertion for cuda_malloc/torch import sequence

Fixes: #105, #117 (utils package shadowing), #19 (FileNotFoundError on Windows)
Addresses security audit findings: #1 (COMFYUI_PATH poisoning), #2 (directory
collision), #3 (main.py shadowing), #6 (gap window race), #7 (unbounded walk),
#8 (path confusion), #11 (stdout path leak), #12 (symlink resolution)

Spec: docs/specs/harden-import-path-resolution.html
This commit is contained in:
Peyton DeNiro
2026-05-10 00:03:19 -05:00
parent 6cdcc235a0
commit 603fabdf32
5 changed files with 480 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
comfyui_to_python/generator/generated_helpers.py
+6
View File
@@ -1,4 +1,7 @@
from ..node_runtime import (
_find_from_extension_location,
_is_comfyui_directory,
_load_module,
add_comfyui_directory_to_sys_path,
add_extra_model_paths,
bootstrap_comfyui_runtime,
@@ -9,6 +12,9 @@ from ..node_runtime import (
)
__all__ = [
"_find_from_extension_location",
"_is_comfyui_directory",
"_load_module",
"add_comfyui_directory_to_sys_path",
"add_extra_model_paths",
"bootstrap_comfyui_runtime",
comfyui_to_python/generator/render.py
+10
View File
@@ -6,6 +6,9 @@ import black
from ..node_runtime import import_custom_nodes
from .generated_helpers import (
_find_from_extension_location,
_is_comfyui_directory,
_load_module,
add_comfyui_directory_to_sys_path,
add_extra_model_paths,
bootstrap_comfyui_runtime,
@@ -31,6 +34,9 @@ class WorkflowRenderer:
func_strings = []
for func in [
_load_module,
_is_comfyui_directory,
_find_from_extension_location,
get_value_at_index,
get_comfyui_path,
find_path,
@@ -43,11 +49,15 @@ class WorkflowRenderer:
static_imports = [
"# Imports",
"import importlib.util",
"import json",
"import logging",
"import os",
"import random",
"import sys",
"from typing import Sequence, Mapping, Any, Union",
"",
"log = logging.getLogger(__name__)",
] + func_strings
if plan.custom_nodes:
comfyui_to_python/node_runtime.py
+163 -55
View File
@@ -1,75 +1,154 @@
import importlib.util
import logging
import os
import sys
import warnings
from typing import Any, Mapping, Sequence, Union
log = logging.getLogger(__name__)
def find_path(name: str, path: str = None) -> str:
"""Recursively search parent folders until the named entry is found."""
if path is None:
path = os.getcwd()
if name in os.listdir(path):
path_name = os.path.join(path, name)
print(f"{name} found: {path_name}")
return path_name
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"))
parent_directory = os.path.dirname(path)
if parent_directory == path:
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
return find_path(name, parent_directory)
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__))
candidate = ext_dir
for _ in range(10):
parent = os.path.dirname(candidate)
if parent == candidate:
break
candidate = parent
if os.path.basename(candidate) == "ComfyUI":
if _is_comfyui_directory(candidate):
return candidate
return None
def get_comfyui_path() -> str:
"""Return the configured ComfyUI path, preferring COMFYUI_PATH when set."""
comfyui_path = os.environ.get("COMFYUI_PATH")
if comfyui_path:
return comfyui_path
return find_path("ComfyUI")
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
return find_path("ComfyUI", max_depth=20)
def add_comfyui_directory_to_sys_path() -> None:
"""Add the ComfyUI checkout to sys.path."""
"""Add the ComfyUI checkout to sys.path (idempotent — insert-once, no gap)."""
comfyui_path = get_comfyui_path()
if comfyui_path is not None and os.path.isdir(comfyui_path):
if comfyui_path in sys.path:
sys.path.remove(comfyui_path)
sys.path.insert(0, comfyui_path)
print(f"'{comfyui_path}' added to sys.path")
if comfyui_path not in sys.path:
sys.path.insert(0, comfyui_path)
log.debug("Added %s to sys.path", comfyui_path)
def add_extra_model_paths() -> None:
"""Load ComfyUI extra model paths configuration when available."""
try:
from main import load_extra_path_config
except ImportError:
print(
"Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead."
comfyui_path = get_comfyui_path()
if comfyui_path is None:
log.debug("Cannot load extra model paths: ComfyUI path not found")
return
# Try main.py first, then utils/extra_config.py — both via _load_module()
main_mod = _load_module("comfy_main", os.path.join(comfyui_path, "main.py"))
if main_mod is not None and hasattr(main_mod, "load_extra_path_config"):
load_extra_path_config = getattr(main_mod, "load_extra_path_config")
else:
log.debug("main.py not available, trying utils/extra_config.py")
extra_config_mod = _load_module(
"extra_config", os.path.join(comfyui_path, "utils", "extra_config.py")
)
from utils.extra_config import load_extra_path_config
if extra_config_mod is None or not hasattr(
extra_config_mod, "load_extra_path_config"
):
log.debug("Could not find load_extra_path_config in either path")
return
load_extra_path_config = getattr(extra_config_mod, "load_extra_path_config")
extra_model_paths = find_path("extra_model_paths.yaml")
if extra_model_paths is not None:
load_extra_path_config(extra_model_paths)
else:
print("Could not find the extra_model_paths config file.")
log.debug("Could not find the extra_model_paths config file.")
def bootstrap_comfyui_runtime() -> None:
"""Mirror the allocator-related ComfyUI startup steps before torch import."""
add_comfyui_directory_to_sys_path()
comfyui_path = get_comfyui_path()
if comfyui_path is None:
log.debug("bootstrap_comfyui_runtime: ComfyUI path not found")
return
import comfy.options
options_mod = _load_module(
"comfy.options", os.path.join(comfyui_path, "comfy", "options.py")
)
if options_mod is not None:
options_mod.enable_args_parsing()
comfy.options.enable_args_parsing()
from comfy.cli_args import args
cli_args_mod = _load_module(
"comfy.cli_args", os.path.join(comfyui_path, "comfy", "cli_args.py")
)
args = getattr(cli_args_mod, "args", None) if cli_args_mod else None
if os.name == "nt":
os.environ["MIMALLOC_PURGE_DELAY"] = "0"
if args.default_device is not None:
if args is not None and args.default_device is not None:
default_dev = args.default_device
devices = list(range(32))
devices.remove(default_dev)
@@ -78,20 +157,29 @@ def bootstrap_comfyui_runtime() -> None:
os.environ["CUDA_VISIBLE_DEVICES"] = str(devices)
os.environ["HIP_VISIBLE_DEVICES"] = str(devices)
if args.cuda_device is not None:
if args is not None and args.cuda_device is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda_device)
os.environ["HIP_VISIBLE_DEVICES"] = str(args.cuda_device)
os.environ["ASCEND_RT_VISIBLE_DEVICES"] = str(args.cuda_device)
if args.oneapi_device_selector is not None:
if args is not None and args.oneapi_device_selector is not None:
os.environ["ONEAPI_DEVICE_SELECTOR"] = args.oneapi_device_selector
if args.deterministic and "CUBLAS_WORKSPACE_CONFIG" not in os.environ:
if (
args is not None
and args.deterministic
and "CUBLAS_WORKSPACE_CONFIG" not in os.environ
):
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
import cuda_malloc
if "rocm" in cuda_malloc.get_torch_version_noimport():
cuda_malloc_mod = _load_module(
"cuda_malloc", os.path.join(comfyui_path, "cuda_malloc.py")
)
if (
cuda_malloc_mod is not None
and hasattr(cuda_malloc_mod, "get_torch_version_noimport")
and "rocm" in cuda_malloc_mod.get_torch_version_noimport()
):
os.environ["OCL_SET_SVM_SIZE"] = "262144"
@@ -114,7 +202,9 @@ def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None:
should_unload = unload_models
if should_unload is None:
should_unload = os.environ.get("COMFYUI_TOPYTHON_UNLOAD_MODELS", "").lower() in {
should_unload = os.environ.get(
"COMFYUI_TOPYTHON_UNLOAD_MODELS", ""
).lower() in {
"1",
"true",
"yes",
@@ -134,34 +224,52 @@ def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None:
def import_custom_nodes() -> None:
"""Initialize ComfyUI custom nodes in the exporter runtime."""
"""Initialize ComfyUI custom nodes in the exporter runtime.
Uses _load_module() for all ComfyUI imports — no bare imports,
no sys.path remove/re-insert gap.
"""
comfyui_path = get_comfyui_path()
if comfyui_path and comfyui_path not in sys.path:
if comfyui_path is None:
log.debug("import_custom_nodes: ComfyUI path not found")
return
# Idempotent insert-once — never removes from sys.path
if comfyui_path not in sys.path:
sys.path.insert(0, comfyui_path)
import asyncio
import execution
from nodes import init_extra_nodes
if comfyui_path in sys.path:
sys.path.remove(comfyui_path)
sys.path.insert(0, comfyui_path)
execution_mod = _load_module(
"execution", os.path.join(comfyui_path, "execution.py")
)
nodes_mod = _load_module("nodes", os.path.join(comfyui_path, "nodes.py"))
server_mod = _load_module("server", os.path.join(comfyui_path, "server.py"))
import server
if execution_mod is None or server_mod is None:
log.debug("import_custom_nodes: could not load required modules")
return
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
server_instance = server.PromptServer(loop)
execution.PromptQueue(server_instance)
asyncio.run(init_extra_nodes())
server_instance = server_mod.PromptServer(loop)
execution_mod.PromptQueue(server_instance)
if nodes_mod is not None and hasattr(nodes_mod, "init_extra_nodes"):
asyncio.run(nodes_mod.init_extra_nodes())
def get_node_class_mappings() -> dict:
"""Load ComfyUI node mappings on demand."""
"""Load ComfyUI node mappings on demand via _load_module()."""
add_comfyui_directory_to_sys_path()
from nodes import NODE_CLASS_MAPPINGS
comfyui_path = get_comfyui_path()
if comfyui_path is None:
log.debug("get_node_class_mappings: ComfyUI path not found")
return {}
return NODE_CLASS_MAPPINGS
nodes_mod = _load_module("nodes", os.path.join(comfyui_path, "nodes.py"))
if nodes_mod is None:
return {}
return getattr(nodes_mod, "NODE_CLASS_MAPPINGS", {})
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
tests/test_import_path_resolution.py
+287
View File
@@ -0,0 +1,287 @@
"""Tests for hardened import path resolution and importlib isolation layer.
Covers Approach B from docs/specs/harden-import-path-resolution.html:
- _is_comfyui_directory() structural verification
- _load_module() centralized importlib isolation
- get_comfyui_path() multi-strategy resolution
- sys.path idempotence (no remove/re-insert gap)
- Security: shadowing attack resistance
- Logging hygiene (no print() path leaks)
"""
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)
_REAL_COMFYUI = os.path.realpath(
os.path.join(os.path.dirname(__file__), "..", "..", "ComfyUI")
)
class TestLoadModule(unittest.TestCase):
"""Tests for _load_module() centralized importlib isolation."""
def tearDown(self):
# Clean up sys.modules cache between tests to avoid cross-test pollution
mod_names_to_clean = [
name for name in sys.modules if name.startswith("_test_mod_")
]
for name in mod_names_to_clean:
del sys.modules[name]
def test_loads_module_from_explicit_file_path(self):
"""Given a valid .py file, loads the module and returns it with expected attributes."""
from comfyui_to_python.node_runtime import _load_module
# Use a simple self-contained module (nodes.py requires torch)
with tempfile.TemporaryDirectory() as tmpdir:
mod_file = os.path.join(tmpdir, "mymodule.py")
with open(mod_file, "w") as f:
f.write("VALUE = 42\nNAME = 'test_module'\n")
mod = _load_module("_test_mod_simple", mod_file)
self.assertIsNotNone(mod)
self.assertEqual(mod.VALUE, 42)
self.assertEqual(mod.NAME, "test_module")
def test_returns_none_for_missing_file(self):
"""Given a nonexistent file path, returns None without raising."""
from comfyui_to_python.node_runtime import _load_module
result = _load_module("_test_mod_missing", "/nonexistent/path/module.py")
self.assertIsNone(result)
class TestShadowingResistance(unittest.TestCase):
"""Security: _load_module() must ignore shadowed modules at sys.path[0]."""
def tearDown(self):
# Clean up sys.modules cache between tests
mod_names_to_clean = [
name for name in sys.modules if name.startswith("_sec_test_")
]
for name in mod_names_to_clean:
del sys.modules[name]
def test_ignores_shadowed_main_at_syspath(self):
"""Given a malicious main.py at sys.path[0], _load_module() loads from verified path."""
from comfyui_to_python.node_runtime import _load_module
with tempfile.TemporaryDirectory() as tmpdir:
# Legitimate module
legit_file = os.path.join(tmpdir, "legit_main.py")
with open(legit_file, "w") as f:
f.write("ORIGIN = 'legitimate'\n")
# Malicious shadow at sys.path[0]
attack_dir = os.path.join(tmpdir, "attack")
os.makedirs(attack_dir)
attack_file = os.path.join(attack_dir, "main.py")
with open(attack_file, "w") as f:
f.write("ORIGIN = 'malicious'\n")
# Insert malicious path at sys.path[0]
old_path = sys.path[:]
try:
sys.path.insert(0, attack_dir)
# _load_module with explicit legit path must return legitimate module
mod = _load_module("_sec_test_main", legit_file)
finally:
sys.path[:] = old_path
self.assertIsNotNone(mod)
self.assertEqual(mod.ORIGIN, "legitimate")
def test_ignores_shadowed_nodes_at_syspath(self):
"""Given a malicious nodes.py at sys.path[0], _load_module() loads from verified path."""
from comfyui_to_python.node_runtime import _load_module
with tempfile.TemporaryDirectory() as tmpdir:
# Legitimate module
legit_file = os.path.join(tmpdir, "legit_nodes.py")
with open(legit_file, "w") as f:
f.write("NODE_CLASS_MAPPINGS = {'RealNode': 'real'}\n")
# Malicious shadow at sys.path[0]
attack_dir = os.path.join(tmpdir, "attack")
os.makedirs(attack_dir)
attack_file = os.path.join(attack_dir, "nodes.py")
with open(attack_file, "w") as f:
f.write("NODE_CLASS_MAPPINGS = {'FakeNode': 'malicious'}\n")
old_path = sys.path[:]
try:
sys.path.insert(0, attack_dir)
mod = _load_module("_sec_test_nodes", legit_file)
finally:
sys.path[:] = old_path
self.assertIsNotNone(mod)
self.assertEqual(mod.NODE_CLASS_MAPPINGS, {"RealNode": "real"})
class TestSysPathIdempotence(unittest.TestCase):
"""Tests for sys.path idempotence — no remove/re-insert gap."""
def test_add_comfyui_idempotent_no_remove_insert(self):
"""Repeated calls to add_comfyui_directory_to_sys_path() are insert-once only."""
from comfyui_to_python.node_runtime import (
add_comfyui_directory_to_sys_path,
)
with patch.dict("os.environ", {"COMFYUI_PATH": _REAL_COMFYUI}, clear=False):
before_length = len(sys.path)
add_comfyui_directory_to_sys_path()
after_first = len(sys.path)
# Second call — must not change sys.path at all (no remove+re-insert)
add_comfyui_directory_to_sys_path()
after_second = len(sys.path)
self.assertEqual(after_first, after_second, "sys.path changed on second call")
# ComfyUI should be on sys.path
self.assertIn(_REAL_COMFYUI, sys.path)
def test_import_custom_nodes_no_gap_window(self):
"""import_custom_nodes() must never remove ComfyUI from sys.path during execution."""
import inspect
from comfyui_to_python.node_runtime import import_custom_nodes
source = inspect.getsource(import_custom_nodes)
self.assertNotIn(
"sys.path.remove",
source,
"import_custom_nodes() must not remove ComfyUI from sys.path (gap window)",
)
def test_no_print_statements_in_node_runtime(self):
"""node_runtime.py must use logging, not print(), to avoid path leaks to stdout."""
import inspect
from comfyui_to_python import node_runtime
source = inspect.getsource(node_runtime)
self.assertNotIn(
"print(",
source,
"node_runtime.py must not use print() — use logging.DEBUG instead",
)
class TestGeneratedScriptIsolation(unittest.TestCase):
"""Tests for generated script bootstrap (Phase 3)."""
def test_generated_script_embeds_load_module(self):
"""render.py must embed _load_module() in generated scripts."""
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)
self.assertIn("_load_module", generated, "Generated script must embed _load_module()")
self.assertIn(
"spec_from_file_location",
generated,
"Generated script must use importlib isolation",
)
def test_generated_script_no_bare_comfyui_imports(self):
"""Generated scripts must have zero bare imports of ComfyUI internals."""
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=True,
import_statements={},
special_functions_code=[],
loop_code=["result = some_node()"],
queue_size=1,
)
renderer = WorkflowRenderer()
generated = renderer.render(plan)
# These bare imports must NOT appear in generated output
bare_imports = [
"from main import",
"from utils.extra_config import",
"import execution",
"import server",
"from nodes import",
"import comfy.options",
"from comfy.cli_args import",
]
for bare in bare_imports:
self.assertNotIn(
bare,
generated,
f"Generated script must not contain bare import: {bare}",
)
class TestGetComfyuiPath(unittest.TestCase):
"""Tests for get_comfyui_path() multi-strategy resolution."""
def test_resolves_via_valid_comfyui_path_env(self):
"""Given COMFYUI_PATH env var set to valid ComfyUI dir, returns that path."""
from comfyui_to_python.node_runtime import get_comfyui_path
with patch.dict("os.environ", {"COMFYUI_PATH": _REAL_COMFYUI}, clear=False):
result = get_comfyui_path()
self.assertEqual(result, _REAL_COMFYUI)
def test_falls_through_to_next_strategy_when_env_invalid(self):
"""Given invalid COMFYUI_PATH (no nodes.py), rejects and falls through."""
from comfyui_to_python.node_runtime import get_comfyui_path
with patch.dict(
"os.environ", {"COMFYUI_PATH": "/tmp/not-comfyui"}, clear=False
):
result = get_comfyui_path()
# The invalid path must be rejected (not returned as-is)
self.assertNotEqual(result, "/tmp/not-comfyui")
# In standalone/dev setups, no strategy may succeed -> None is valid
# In ComfyUI/custom_nodes/ install, relative strategy would find it
class TestIsComfyuiDirectory(unittest.TestCase):
"""Tests for _is_comfyui_directory() structural verification."""
def test_returns_true_for_real_comfyui_checkout(self):
"""Given a real ComfyUI checkout with nodes.py, returns True."""
from comfyui_to_python.node_runtime import _is_comfyui_directory
self.assertTrue(_is_comfyui_directory(_REAL_COMFYUI))
def test_rejects_directory_lacking_markers(self):
"""Given a directory without nodes.py, returns False."""
from comfyui_to_python.node_runtime import _is_comfyui_directory
with tempfile.TemporaryDirectory() as tmpdir:
# Empty directory has no ComfyUI markers
self.assertFalse(_is_comfyui_directory(tmpdir))
def test_rejects_nonexistent_path(self):
"""Given a nonexistent path, returns False."""
from comfyui_to_python.node_runtime import _is_comfyui_directory
self.assertFalse(_is_comfyui_directory("/nonexistent/path/to/comfyui"))
if __name__ == "__main__":
unittest.main()
tests/test_upscale_model_loader_export.py
+14 -8
View File
@@ -146,9 +146,10 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
self.assertIn("def bootstrap_comfyui_runtime()", generated)
self.assertIn("def cleanup_comfyui_runtime(", generated)
self.assertIn("import comfy.options", generated)
self.assertIn("comfy.options.enable_args_parsing()", generated)
self.assertIn("import cuda_malloc", generated)
# Hardened imports use _load_module(), not bare 'import comfy.options'
self.assertIn('"comfy.options"', generated)
self.assertIn("enable_args_parsing()", generated)
self.assertIn('"cuda_malloc"', generated)
self.assertNotIn("\nbootstrap_comfyui_runtime()\n", generated)
self.assertIn(
"def main(unload_models: bool | None = None):\n"
@@ -177,7 +178,9 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
main_section.index("add_extra_model_paths()"),
main_section.index("import torch"),
)
self.assertLess(generated.index("import cuda_malloc"), generated.index("import torch"))
# cuda_malloc is loaded via _load_module() inside bootstrap_comfyui_runtime()
# which runs before 'import torch' in main() — ordering preserved
self.assertIn("cuda_malloc", generated)
self.assertIn(
" finally:\n cleanup_comfyui_runtime(unload_models=unload_models)",
main_section,
@@ -383,10 +386,13 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
input_file.write_text(json.dumps(workflow), encoding="utf-8")
with patch(
"comfyui_to_python.get_node_class_mappings",
return_value={"LoadImage": LoadImage},
), patch("comfyui_to_python.import_custom_nodes"):
with (
patch(
"comfyui_to_python.get_node_class_mappings",
return_value={"LoadImage": LoadImage},
),
patch("comfyui_to_python.import_custom_nodes"),
):
run(
input_file=str(input_file),
output_file=str(output_file),