352 lines
11 KiB
Python
352 lines
11 KiB
Python
"""Tests for the SpectrumifyNode ComfyUI integration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
import nodes as node_mod
|
|
from nodes import (
|
|
COLOR_MODE_LABELS,
|
|
SIZE_LABELS,
|
|
SpectrumifyNode,
|
|
_random_suffix,
|
|
_save_preview,
|
|
)
|
|
from spectrumify.palette import PALETTE_16, PALETTE_BW, PALETTE_GRAY
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Schema (INPUT_TYPES / RETURN_TYPES)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_input_types_has_required_fields():
|
|
schema = SpectrumifyNode.INPUT_TYPES()
|
|
assert "required" in schema
|
|
required = schema["required"]
|
|
for key in ("image", "size", "color_mode", "pixel_scale"):
|
|
assert key in required
|
|
|
|
|
|
def test_input_types_size_options_match_labels():
|
|
schema = SpectrumifyNode.INPUT_TYPES()
|
|
size_options = schema["required"]["size"][0]
|
|
assert set(size_options) == set(SIZE_LABELS.keys())
|
|
|
|
|
|
def test_input_types_color_mode_options_match_labels():
|
|
schema = SpectrumifyNode.INPUT_TYPES()
|
|
options = schema["required"]["color_mode"][0]
|
|
assert set(options) == set(COLOR_MODE_LABELS.keys())
|
|
|
|
|
|
def test_input_types_pixel_scale_range():
|
|
schema = SpectrumifyNode.INPUT_TYPES()
|
|
spec = schema["required"]["pixel_scale"][1]
|
|
assert spec["min"] == 1
|
|
assert spec["max"] == 16
|
|
assert spec["default"] == 4
|
|
|
|
|
|
def test_return_types_is_image():
|
|
assert SpectrumifyNode.RETURN_TYPES == ("IMAGE",)
|
|
assert SpectrumifyNode.RETURN_NAMES == ("spectrumified",)
|
|
|
|
|
|
def test_function_attribute_is_callable():
|
|
assert SpectrumifyNode.FUNCTION == "spectrumify"
|
|
assert callable(getattr(SpectrumifyNode, "spectrumify"))
|
|
|
|
|
|
def test_category_set():
|
|
assert SpectrumifyNode.CATEGORY == "image/transform"
|
|
|
|
|
|
def test_is_output_node():
|
|
assert SpectrumifyNode.OUTPUT_NODE is True
|
|
|
|
|
|
def test_optional_fields_present():
|
|
schema = SpectrumifyNode.INPUT_TYPES()
|
|
assert "optional" in schema
|
|
optional = schema["optional"]
|
|
assert "save_gif" in optional
|
|
assert "filename_prefix" in optional
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Mappings registered
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_node_class_mappings_contains_node():
|
|
# ComfyUI loads the top-level __init__.py as the package entry point. We
|
|
# load it as a package here so the relative ``from .nodes import ...``
|
|
# inside __init__.py resolves. The module must be registered in
|
|
# sys.modules *before* executing so relative imports can find its parent.
|
|
import importlib.util
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
pkg_name = "spectrumify_comfy_pkg"
|
|
repo_root = Path(__file__).resolve().parents[1]
|
|
spec = importlib.util.spec_from_file_location(
|
|
pkg_name,
|
|
repo_root / "__init__.py",
|
|
submodule_search_locations=[str(repo_root)],
|
|
)
|
|
assert spec is not None and spec.loader is not None
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[pkg_name] = module
|
|
try:
|
|
spec.loader.exec_module(module)
|
|
assert "SpectrumifyNode" in module.NODE_CLASS_MAPPINGS
|
|
assert "SpectrumifyNode" in module.NODE_DISPLAY_NAME_MAPPINGS
|
|
assert (
|
|
module.NODE_CLASS_MAPPINGS["SpectrumifyNode"].__name__
|
|
== SpectrumifyNode.__name__
|
|
)
|
|
finally:
|
|
sys.modules.pop(pkg_name, None)
|
|
sys.modules.pop(f"{pkg_name}.nodes", None)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Execution
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _all_pixels_in_palette(arr: np.ndarray, palette) -> bool:
|
|
unique = {tuple(p) for p in arr.reshape(-1, 3).tolist()}
|
|
return unique.issubset({tuple(c) for c in palette})
|
|
|
|
|
|
def test_spectrumify_runs_with_numpy_tensor(patched_comfy_dirs):
|
|
node = SpectrumifyNode()
|
|
tensor = np.random.default_rng(1).random((1, 64, 64, 3), dtype=np.float32)
|
|
result = node.spectrumify(
|
|
image=tensor,
|
|
size="SMALL (80px)",
|
|
color_mode="16 Colors",
|
|
pixel_scale=2,
|
|
)
|
|
assert "result" in result
|
|
assert "ui" in result
|
|
(out_tensor,) = result["result"]
|
|
arr = np.asarray(out_tensor)
|
|
assert arr.ndim == 4
|
|
assert arr.shape[0] == 1
|
|
assert arr.shape[-1] == 3
|
|
assert arr.min() >= 0.0 and arr.max() <= 1.0
|
|
|
|
|
|
def test_spectrumify_preview_files_written(patched_comfy_dirs):
|
|
node = SpectrumifyNode()
|
|
tensor = np.zeros((1, 20, 20, 3), dtype=np.float32)
|
|
result = node.spectrumify(
|
|
image=tensor,
|
|
size="SMALL (80px)",
|
|
color_mode="Grayscale",
|
|
pixel_scale=1,
|
|
filename_prefix="TestPreview",
|
|
)
|
|
images = result["ui"]["images"]
|
|
assert len(images) == 1
|
|
entry = images[0]
|
|
assert entry["type"] == "temp"
|
|
expected_path = os.path.join(patched_comfy_dirs["temp"], entry["filename"])
|
|
assert os.path.exists(expected_path)
|
|
assert entry["filename"].startswith("TestPreview_")
|
|
|
|
|
|
def test_spectrumify_batch_b_gt_1(patched_comfy_dirs):
|
|
node = SpectrumifyNode()
|
|
tensor = np.random.default_rng(2).random((3, 32, 32, 3), dtype=np.float32)
|
|
result = node.spectrumify(
|
|
image=tensor,
|
|
size="SMALL (80px)",
|
|
color_mode="16 Colors",
|
|
pixel_scale=1,
|
|
)
|
|
(out_tensor,) = result["result"]
|
|
arr = np.asarray(out_tensor)
|
|
assert arr.shape[0] == 3
|
|
assert len(result["ui"]["images"]) == 3
|
|
|
|
|
|
@pytest.mark.parametrize("size", list(SIZE_LABELS.keys()))
|
|
def test_spectrumify_all_sizes(patched_comfy_dirs, size):
|
|
node = SpectrumifyNode()
|
|
tensor = np.random.default_rng(3).random((1, 32, 32, 3), dtype=np.float32)
|
|
result = node.spectrumify(
|
|
image=tensor,
|
|
size=size,
|
|
color_mode="16 Colors",
|
|
pixel_scale=1,
|
|
)
|
|
(out_tensor,) = result["result"]
|
|
arr = np.asarray(out_tensor)
|
|
assert arr.shape[2] == SIZE_LABELS[size]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"color_mode, palette",
|
|
[
|
|
("16 Colors", PALETTE_16),
|
|
("Grayscale", PALETTE_GRAY),
|
|
("Black & White", PALETTE_BW),
|
|
],
|
|
)
|
|
def test_spectrumify_all_color_modes(patched_comfy_dirs, color_mode, palette):
|
|
node = SpectrumifyNode()
|
|
tensor = np.random.default_rng(4).random((1, 24, 24, 3), dtype=np.float32)
|
|
result = node.spectrumify(
|
|
image=tensor,
|
|
size="SMALL (80px)",
|
|
color_mode=color_mode,
|
|
pixel_scale=1,
|
|
)
|
|
(out_tensor,) = result["result"]
|
|
pixels = (np.asarray(out_tensor)[0] * 255.0).clip(0, 255).astype(np.uint8)
|
|
assert _all_pixels_in_palette(pixels, palette)
|
|
|
|
|
|
@pytest.mark.parametrize("size", list(SIZE_LABELS.keys()))
|
|
@pytest.mark.parametrize("color_mode", list(COLOR_MODE_LABELS.keys()))
|
|
def test_spectrumify_all_9_combinations(patched_comfy_dirs, size, color_mode):
|
|
node = SpectrumifyNode()
|
|
tensor = np.random.default_rng(5).random((1, 30, 30, 3), dtype=np.float32)
|
|
result = node.spectrumify(
|
|
image=tensor,
|
|
size=size,
|
|
color_mode=color_mode,
|
|
pixel_scale=1,
|
|
)
|
|
(out_tensor,) = result["result"]
|
|
assert np.asarray(out_tensor).shape[2] == SIZE_LABELS[size]
|
|
|
|
|
|
def test_spectrumify_save_gif_creates_file(patched_comfy_dirs):
|
|
node = SpectrumifyNode()
|
|
tensor = np.random.default_rng(6).random((2, 24, 24, 3), dtype=np.float32)
|
|
node.spectrumify(
|
|
image=tensor,
|
|
size="SMALL (80px)",
|
|
color_mode="16 Colors",
|
|
pixel_scale=1,
|
|
save_gif=True,
|
|
filename_prefix="GifTest",
|
|
)
|
|
gifs = [
|
|
f for f in os.listdir(patched_comfy_dirs["output"]) if f.lower().endswith(".gif")
|
|
]
|
|
assert gifs, "expected at least one GIF to be written"
|
|
|
|
|
|
def test_spectrumify_rejects_unknown_size():
|
|
node = SpectrumifyNode()
|
|
tensor = np.zeros((1, 4, 4, 3), dtype=np.float32)
|
|
with pytest.raises(ValueError):
|
|
node.spectrumify(
|
|
image=tensor,
|
|
size="HUGE (9000px)",
|
|
color_mode="16 Colors",
|
|
pixel_scale=1,
|
|
)
|
|
|
|
|
|
def test_spectrumify_rejects_unknown_color_mode():
|
|
node = SpectrumifyNode()
|
|
tensor = np.zeros((1, 4, 4, 3), dtype=np.float32)
|
|
with pytest.raises(ValueError):
|
|
node.spectrumify(
|
|
image=tensor,
|
|
size="SMALL (80px)",
|
|
color_mode="Cyberpunk",
|
|
pixel_scale=1,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_random_suffix_length_and_alphabet():
|
|
suffix = _random_suffix(8)
|
|
assert len(suffix) == 8
|
|
assert all(c in "0123456789abcdef" for c in suffix)
|
|
|
|
|
|
def test_save_preview_sanitises_prefix(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(node_mod, "_temp_directory", lambda: str(tmp_path))
|
|
img = Image.new("RGB", (4, 4), (1, 2, 3))
|
|
descriptor = _save_preview(img, "Bad/Prefix!!", 0)
|
|
assert descriptor["type"] == "temp"
|
|
assert "/" not in descriptor["filename"]
|
|
assert os.path.exists(os.path.join(tmp_path, descriptor["filename"]))
|
|
|
|
|
|
def test_save_preview_handles_empty_prefix(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(node_mod, "_temp_directory", lambda: str(tmp_path))
|
|
img = Image.new("RGB", (2, 2), (0, 0, 0))
|
|
descriptor = _save_preview(img, "", 3)
|
|
assert descriptor["filename"].startswith("Spectrumify_")
|
|
|
|
|
|
def test_folder_paths_module_returns_none_without_comfy(monkeypatch):
|
|
"""When ComfyUI is not installed the helper must return ``None``."""
|
|
|
|
import sys
|
|
|
|
monkeypatch.setitem(sys.modules, "folder_paths", None)
|
|
# monkeypatching to None makes ``import folder_paths`` raise ImportError,
|
|
# which our helper catches and converts to None.
|
|
assert node_mod._folder_paths_module() is None
|
|
|
|
|
|
def test_temp_directory_fallback_creates_local_dir(monkeypatch):
|
|
"""Fallback temp dir is a ``temp/`` folder next to nodes.py."""
|
|
|
|
monkeypatch.setattr(node_mod, "_folder_paths_module", lambda: None)
|
|
path = node_mod._temp_directory()
|
|
assert os.path.isdir(path)
|
|
assert os.path.basename(path) == "temp"
|
|
|
|
|
|
def test_output_directory_fallback_creates_local_dir(monkeypatch):
|
|
"""Fallback output dir is an ``output/`` folder next to nodes.py."""
|
|
|
|
monkeypatch.setattr(node_mod, "_folder_paths_module", lambda: None)
|
|
path = node_mod._output_directory()
|
|
assert os.path.isdir(path)
|
|
assert os.path.basename(path) == "output"
|
|
|
|
|
|
def test_temp_directory_uses_comfy_folder_paths(tmp_path, monkeypatch):
|
|
"""When ``folder_paths`` is available, the helper uses ``get_temp_directory``."""
|
|
|
|
class FakeFolderPaths:
|
|
@staticmethod
|
|
def get_temp_directory():
|
|
return str(tmp_path / "comfy_temp")
|
|
|
|
monkeypatch.setattr(node_mod, "_folder_paths_module", lambda: FakeFolderPaths)
|
|
path = node_mod._temp_directory()
|
|
assert path == str(tmp_path / "comfy_temp")
|
|
assert os.path.isdir(path)
|
|
|
|
|
|
def test_output_directory_uses_comfy_folder_paths(tmp_path, monkeypatch):
|
|
"""When ``folder_paths`` is available, the helper uses ``get_output_directory``."""
|
|
|
|
class FakeFolderPaths:
|
|
@staticmethod
|
|
def get_output_directory():
|
|
return str(tmp_path / "comfy_output")
|
|
|
|
monkeypatch.setattr(node_mod, "_folder_paths_module", lambda: FakeFolderPaths)
|
|
path = node_mod._output_directory()
|
|
assert path == str(tmp_path / "comfy_output")
|
|
assert os.path.isdir(path)
|