S1 of the 2026-09-08 layout plan: DyPE/SEGA/SPA/HAP node classes now live
in nodes/{dype,sega,spa,hap}.py with engine imports via dual-form try/except
(loader-relative ..src / flat src). Entry __init__.py keeps extension
registration only. nodes/__init__.py re-exports the full 8-class surface
(transitional src/*_node re-exports until S2/S3 land). All patch nodes
category -> WMNodes/image. New tests/test_structure.py guards the layout.
94 lines
3.4 KiB
Python
94 lines
3.4 KiB
Python
"""Tests for the SEGA ComfyUI node definition (Tier 2: node schema tests)."""
|
|
import pathlib
|
|
|
|
import pytest
|
|
|
|
_ROOT = pathlib.Path(__file__).parent.parent
|
|
_ENTRY = _ROOT / "__init__.py"
|
|
_NODE = _ROOT / "nodes" / "sega.py"
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSegaNodeSchema:
|
|
def _get_node_class(self):
|
|
"""Import the SEGA node class from the pack entry."""
|
|
import importlib
|
|
spec = importlib.util.spec_from_file_location(
|
|
"dype_init", pathlib.Path(__file__).parent.parent / "__init__.py"
|
|
)
|
|
mod = importlib.util.module_from_spec(spec)
|
|
# We need comfy_api which may not be available in test env
|
|
try:
|
|
spec.loader.exec_module(mod)
|
|
return mod.SEGA
|
|
except Exception:
|
|
pytest.skip("comfy_api not available in test environment")
|
|
|
|
def test_node_class_exists(self):
|
|
"""SEGA node class should be defined in nodes/sega.py."""
|
|
content = _NODE.read_text(encoding="utf-8")
|
|
assert "class SEGA" in content
|
|
|
|
def test_node_registered_in_extension(self):
|
|
"""SEGA should be in the extension's node list."""
|
|
content = _ENTRY.read_text(encoding="utf-8")
|
|
assert "SEGA" in content
|
|
# Check that get_node_list includes SEGA
|
|
assert "return [DyPE_FLUX, SEGA" in content
|
|
|
|
def test_imports_apply_sega(self):
|
|
"""nodes/sega.py should import apply_sega_to_model."""
|
|
content = _NODE.read_text(encoding="utf-8")
|
|
assert "apply_sega_to_model" in content
|
|
|
|
def test_node_has_sega_inputs(self):
|
|
"""The SEGA node definition should include SEGA-specific inputs."""
|
|
content = _NODE.read_text(encoding="utf-8")
|
|
sega_inputs = [
|
|
"mscale_alpha",
|
|
"mscale_beta",
|
|
"mscale_min",
|
|
"spread_min",
|
|
"spread_max",
|
|
"spread_alpha",
|
|
"base_mscale_formula",
|
|
"base_mscale_coefficient",
|
|
]
|
|
for inp in sega_inputs:
|
|
assert inp in content, f"SEGA node should have input: {inp}"
|
|
|
|
def test_node_has_method_combo(self):
|
|
"""The SEGA node should have a method combo with 'sega' and 'ntk' options."""
|
|
content = _NODE.read_text(encoding="utf-8")
|
|
assert '"sega"' in content
|
|
assert '"ntk"' in content
|
|
|
|
def test_node_has_model_type_combo(self):
|
|
"""The SEGA node should have a model_type combo."""
|
|
content = _NODE.read_text(encoding="utf-8")
|
|
assert '"auto"' in content
|
|
assert '"flux"' in content
|
|
assert '"qwen"' in content
|
|
|
|
def test_node_defaults_match_paper(self):
|
|
"""Default SEGA parameters should match the paper's recommendations."""
|
|
content = _NODE.read_text(encoding="utf-8")
|
|
# mscale_alpha = 0.15 (paper default)
|
|
assert "default=0.15" in content
|
|
# mscale_beta = 1.5 (paper default)
|
|
assert "default=1.5" in content
|
|
# base_mscale_coefficient = 0.08 (paper kappa)
|
|
assert "default=0.08" in content
|
|
# spread_alpha = 1.5 (paper gamma)
|
|
assert "default=1.5" in content
|
|
|
|
def test_node_has_output(self):
|
|
"""The SEGA node should output a patched model."""
|
|
content = _NODE.read_text(encoding="utf-8")
|
|
assert "io.Model.Output" in content
|
|
|
|
def test_node_category(self):
|
|
"""The SEGA node should be in the WMNodes/image category."""
|
|
content = _NODE.read_text(encoding="utf-8")
|
|
assert "WMNodes/image" in content
|