glonlas-ComfyUI-image-profile/tests/test_nodes.py

"""Unit tests for nodes.py."""

from __future__ import annotations

import importlib.util
import json
import sys
import types
import unittest
from pathlib import Path


NODES_PATH = Path(__file__).resolve().parents[1] / "nodes.py"


def load_nodes_module():
    """Load nodes.py with mocked external runtime modules."""
    module_name = "nodes_under_test"
    sys.modules.pop(module_name, None)

    fake_torch = types.ModuleType("torch")
    fake_torch.zeros = lambda shape, device=None: {"shape": shape, "device": device}

    fake_comfy = types.ModuleType("comfy")
    fake_model_management = types.ModuleType("comfy.model_management")
    fake_model_management.intermediate_device = lambda: "cpu"
    fake_comfy.model_management = fake_model_management

    previous = {
        "torch": sys.modules.get("torch"),
        "comfy": sys.modules.get("comfy"),
        "comfy.model_management": sys.modules.get("comfy.model_management"),
    }
    sys.modules["torch"] = fake_torch
    sys.modules["comfy"] = fake_comfy
    sys.modules["comfy.model_management"] = fake_model_management

    try:
        spec = importlib.util.spec_from_file_location(module_name, NODES_PATH)
        if spec is None or spec.loader is None:
            raise RuntimeError("Unable to load nodes.py spec")
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        return module
    finally:
        for name, original in previous.items():
            if original is None:
                sys.modules.pop(name, None)
            else:
                sys.modules[name] = original


class TestComfyUIImageProfile(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.nodes = load_nodes_module()
        cls.node_cls = cls.nodes.ComfyUIImageProfile

    def test_sanitize_dimension_rounds_and_clamps(self):
        self.assertEqual(self.node_cls._sanitize_dimension(11), 8)
        self.assertEqual(self.node_cls._sanitize_dimension(12), 16)
        self.assertEqual(self.node_cls._sanitize_dimension(-100), 8)
        self.assertEqual(self.node_cls._sanitize_dimension(20000), 16384)

    def test_load_profiles_falls_back_to_defaults_on_invalid_json(self):
        profiles = self.node_cls._load_profiles("not-json")
        self.assertEqual(len(profiles), len(self.nodes.DEFAULT_PROFILES))
        self.assertEqual(profiles[0]["id"], "default-landscape-low")

    def test_load_profiles_filters_invalid_profiles(self):
        payload = json.dumps(
            [
                {"id": "ok", "name": "Valid", "width": 512, "height": 768, "steps": 20},
                {"id": "missing-name", "width": 512, "height": 768, "steps": 20},
                {"name": "missing-id", "width": 512, "height": 768, "steps": 20},
            ]
        )
        profiles = self.node_cls._load_profiles(payload)
        self.assertEqual(len(profiles), 1)
        self.assertEqual(profiles[0]["id"], "ok")
        self.assertEqual(profiles[0]["name"], "Valid")
        self.assertEqual(profiles[0]["width"], 512)
        self.assertEqual(profiles[0]["height"], 768)
        self.assertEqual(profiles[0]["steps"], 20)

    def test_resolve_selected_profile_falls_back_to_first_profile(self):
        node = self.node_cls()
        profiles = [
            {"id": "first", "name": "First", "width": 512, "height": 768, "steps": 20},
            {"id": "second", "name": "Second", "width": 768, "height": 512, "steps": 30},
        ]
        selected = node._resolve_selected_profile(
            profiles,
            selected_profile_id="missing",
            selected_width=111,
            selected_height=222,
            selected_steps=0,
        )
        self.assertEqual(selected["id"], "first")
        self.assertEqual(selected["name"], "Selected")
        self.assertEqual(selected["width"], 112)
        self.assertEqual(selected["height"], 224)
        self.assertEqual(selected["steps"], 1)

    def test_generate_returns_latent_and_steps(self):
        node = self.node_cls()
        profiles_json = json.dumps(
            [
                {"id": "p1", "name": "Profile One", "width": 1152, "height": 864, "steps": 14},
            ]
        )
        latent, steps = node.generate(
            profiles_json=profiles_json,
            selected_profile_id="p1",
            selected_width=8,
            selected_height=8,
            selected_steps=1,
        )

        self.assertEqual(steps, 14)
        self.assertEqual(latent["samples"]["shape"], [1, 4, 108, 144])
        self.assertEqual(latent["samples"]["device"], "cpu")

    # -- image-driven aspect ratio -----------------------------------------

    @staticmethod
    def _fake_image(height: int, width: int, batch: int = 1, channels: int = 3):
        # Mimics a ComfyUI IMAGE tensor: shape [batch, height, width, channels].
        return types.SimpleNamespace(shape=(batch, height, width, channels))

    def test_image_dimensions_reads_bhwc_shape(self):
        image = self._fake_image(height=512, width=768)
        self.assertEqual(self.node_cls._image_dimensions(image), (768, 512))

    def test_image_dimensions_reads_hwc_shape(self):
        image = types.SimpleNamespace(shape=(720, 1280, 3))
        self.assertEqual(self.node_cls._image_dimensions(image), (1280, 720))

    def test_image_dimensions_returns_none_for_invalid(self):
        self.assertIsNone(self.node_cls._image_dimensions(None))
        self.assertIsNone(self.node_cls._image_dimensions(types.SimpleNamespace(shape=(3,))))
        self.assertIsNone(self.node_cls._image_dimensions(types.SimpleNamespace(shape=(1, 0, 100, 3))))

    def test_fit_ratio_to_area_keeps_landscape_ratio_and_budget(self):
        # 3:2 image fitted into a 1152x864 (~1.0 MP) profile budget.
        image = self._fake_image(height=512, width=768)
        width, height = self.node_cls._fit_ratio_to_area(image, 1152, 864)
        self.assertGreater(width, height)
        self.assertEqual(width % self.nodes.DIM_MULTIPLE, 0)
        self.assertEqual(height % self.nodes.DIM_MULTIPLE, 0)
        # Roughly preserves the source ratio (1.5) and the pixel budget.
        self.assertAlmostEqual(width / height, 768 / 512, delta=0.05)
        self.assertAlmostEqual((width * height) / (1152 * 864), 1.0, delta=0.05)

    def test_fit_ratio_to_area_produces_portrait_for_tall_image(self):
        image = self._fake_image(height=1024, width=512)
        width, height = self.node_cls._fit_ratio_to_area(image, 1152, 864)
        self.assertLess(width, height)

    def test_fit_ratio_to_area_falls_back_without_shape(self):
        self.assertEqual(self.node_cls._fit_ratio_to_area(None, 1152, 864), (1152, 864))

    def test_generate_uses_image_ratio_with_profile_resolution(self):
        node = self.node_cls()
        profiles_json = json.dumps(
            [{"id": "p1", "name": "Profile One", "width": 1152, "height": 864, "steps": 14}],
        )
        image = self._fake_image(height=512, width=768)
        latent, steps = node.generate(
            profiles_json=profiles_json,
            selected_profile_id="p1",
            selected_width=8,
            selected_height=8,
            selected_steps=1,
            image=image,
        )

        # Steps still come from the profile; latent dims follow the image ratio.
        self.assertEqual(steps, 14)
        _, _, latent_h, latent_w = latent["samples"]["shape"]
        self.assertGreater(latent_w, latent_h)

    def test_generate_without_image_uses_profile_dimensions(self):
        node = self.node_cls()
        profiles_json = json.dumps(
            [{"id": "p1", "name": "Profile One", "width": 1152, "height": 864, "steps": 14}],
        )
        latent, _ = node.generate(
            profiles_json=profiles_json,
            selected_profile_id="p1",
            selected_width=8,
            selected_height=8,
            selected_steps=1,
        )
        # Unchanged behavior: 1152x864 -> latent [1, 4, 108, 144].
        self.assertEqual(latent["samples"]["shape"], [1, 4, 108, 144])


if __name__ == "__main__":
    unittest.main()