"""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)