DumasNodes/tests/test_dumas_image_nodes.py

import importlib
import os
import sys
import tempfile
import types
import unittest


class FakeImageArray:
    def __init__(self, width=8, height=6):
        self.shape = (height, width, 3)

    def cpu(self):
        return self

    def numpy(self):
        return self

    def astype(self, _dtype):
        return self

    def __rmul__(self, _value):
        return self


class FakeTensorBatch:
    def __init__(self, width=8, height=6, count=1):
        self.image = FakeImageArray(width=width, height=height)
        self.shape = (count, height, width, 3)
        self.count = count

    def __getitem__(self, index):
        if isinstance(index, slice):
            return FakeTensorBatch(
                width=self.shape[2],
                height=self.shape[1],
                count=len(range(*index.indices(self.count))),
            )
        if index != 0:
            raise IndexError(index)
        return self.image


class FakePILImage:
    saved_paths = []

    def save(self, path, *args, **kwargs):
        self.saved_paths.append((path, args, kwargs))

    def convert(self, _mode):
        return self


class DumasImageNodeTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.temp_dir = tempfile.mkdtemp(prefix="dumas-image-node-")
        fake_numpy = types.SimpleNamespace(
            clip=lambda array, _low, _high: array,
            uint8="uint8",
        )
        fake_pil_image_module = types.SimpleNamespace(fromarray=lambda _array: FakePILImage())
        fake_pil_module = types.SimpleNamespace(Image=fake_pil_image_module)
        fake_folder_paths = types.SimpleNamespace(
            get_temp_directory=lambda: cls.temp_dir,
            get_output_directory=lambda: cls.temp_dir,
            get_save_image_path=lambda prefix, _out, _width, _height: (
                cls.temp_dir,
                prefix,
                1,
                "",
                prefix,
            ),
        )

        cls._saved_modules = {
            name: sys.modules.get(name)
            for name in ("numpy", "PIL", "PIL.Image", "folder_paths")
        }
        sys.modules["numpy"] = fake_numpy
        sys.modules["PIL"] = fake_pil_module
        sys.modules["PIL.Image"] = fake_pil_image_module
        sys.modules["folder_paths"] = fake_folder_paths
        cls.image_nodes = importlib.import_module("dumas_image_nodes")

    @classmethod
    def tearDownClass(cls):
        for name, module in cls._saved_modules.items():
            if module is None:
                sys.modules.pop(name, None)
            else:
                sys.modules[name] = module

    def setUp(self):
        FakePILImage.saved_paths = []

    def test_compare_images_returns_second_input_as_new_image(self):
        node = self.image_nodes.DumasImageCompareNode()
        image1 = FakeTensorBatch()
        image2 = FakeTensorBatch()

        result = node.compare_images(image1=image1, image2=image2)

        self.assertIs(result["result"][0], image2)
        self.assertEqual([item["slot"] for item in result["ui"]["images"]], [1, 2])
        self.assertEqual(len(FakePILImage.saved_paths), 2)

    def test_compare_images_falls_back_to_first_image(self):
        node = self.image_nodes.DumasImageCompareNode()
        image1 = FakeTensorBatch()

        result = node.compare_images(image1=image1)

        self.assertIs(result["result"][0], image1)
        self.assertEqual([item["slot"] for item in result["ui"]["images"]], [1])

    def test_compare_images_handles_missing_inputs(self):
        node = self.image_nodes.DumasImageCompareNode()

        result = node.compare_images()

        self.assertIsNone(result["result"][0])
        self.assertEqual(result["ui"]["images"], [])

    def test_saved_filenames_use_dumas_prefix(self):
        node = self.image_nodes.DumasImageCompareNode()
        image1 = FakeTensorBatch()

        result = node.compare_images(image1=image1)

        self.assertTrue(result["ui"]["images"][0]["filename"].startswith("dumas_compare"))
        self.assertTrue(os.path.basename(FakePILImage.saved_paths[0][0]).startswith("dumas_compare"))

    def test_save_image_uses_second_input_token(self):
        node = self.image_nodes.DumasSaveImageNode()
        image = FakeTensorBatch(width=10, height=12)

        node.save_images(
            images=image,
            folder=self.temp_dir,
            pattern="shot_%input%_%input2%_%counter%",
            format="png",
            quality=100,
            embed_workflow=False,
            save_on_run=True,
            name="alpha.png",
            name_2="beta/final",
        )

        saved_name = os.path.basename(FakePILImage.saved_paths[0][0])
        self.assertEqual(saved_name, "shot_alpha_beta_final_001.png")

    def test_save_image_returns_ui_entries_for_output_folder(self):
        node = self.image_nodes.DumasSaveImageNode()
        image = FakeTensorBatch()

        result = node.save_images(
            images=image,
            folder="",
            pattern="result_%counter%",
            format="jpg",
            quality=90,
            embed_workflow=False,
            save_on_run=True,
        )

        self.assertEqual(result["ui"]["images"][0]["type"], "output")
        self.assertTrue(result["ui"]["images"][0]["filename"].endswith(".jpg"))

    def test_save_image_returns_ui_entries_for_external_folder(self):
        node = self.image_nodes.DumasSaveImageNode()
        image = FakeTensorBatch()
        external_dir = tempfile.mkdtemp(prefix="dumas-image-node-external-")

        result = node.save_images(
            images=image,
            folder=external_dir,
            pattern="external_%counter%",
            format="png",
            quality=100,
            embed_workflow=False,
            save_on_run=True,
        )

        self.assertEqual(result["ui"]["images"][0]["type"], "external")
        self.assertEqual(result["ui"]["images"][0]["subfolder"], external_dir.replace("\\", "/"))

    def test_save_image_skips_when_save_is_disabled(self):
        node = self.image_nodes.DumasSaveImageNode()
        image = FakeTensorBatch()

        result = node.save_images(
            images=image,
            folder=self.temp_dir,
            pattern="ignored_%counter%",
            format="png",
            quality=100,
            embed_workflow=False,
            save_on_run=False,
        )

        self.assertEqual(result["ui"]["images"], [])
        self.assertEqual(FakePILImage.saved_paths, [])


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