290 lines
9.3 KiB
Python
290 lines
9.3 KiB
Python
import importlib
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import types
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
|
|
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_resolves_date_size_and_batch_tokens(self):
|
|
node = self.image_nodes.DumasSaveImageNode()
|
|
image = FakeTensorBatch(width=10, height=12, count=2)
|
|
fake_now = time.struct_time((2026, 8, 5, 13, 7, 9, 2, 217, -1))
|
|
|
|
with mock.patch.object(self.image_nodes.time, "localtime", return_value=fake_now):
|
|
node.save_images(
|
|
images=image,
|
|
folder=self.temp_dir,
|
|
pattern="asset_%date:yyyy-MM-dd%_%date:hh-mm-ss%_%width%x%height%_%batch_num%_%counter%",
|
|
format="png",
|
|
quality=100,
|
|
embed_workflow=False,
|
|
save_on_run=True,
|
|
)
|
|
|
|
saved_names = [os.path.basename(path) for path, _args, _kwargs in FakePILImage.saved_paths]
|
|
self.assertEqual(
|
|
saved_names,
|
|
[
|
|
"asset_2026-08-05_13-07-09_10x12_0_001.png",
|
|
"asset_2026-08-05_13-07-09_10x12_1_001.png",
|
|
],
|
|
)
|
|
|
|
def test_save_image_counter_increments_for_existing_files(self):
|
|
node = self.image_nodes.DumasSaveImageNode()
|
|
image = FakeTensorBatch()
|
|
|
|
node.save_images(
|
|
images=image,
|
|
folder=self.temp_dir,
|
|
pattern="counter_%counter%",
|
|
format="png",
|
|
quality=100,
|
|
embed_workflow=False,
|
|
save_on_run=True,
|
|
)
|
|
open(FakePILImage.saved_paths[0][0], "a", encoding="utf-8").close()
|
|
node.save_images(
|
|
images=image,
|
|
folder=self.temp_dir,
|
|
pattern="counter_%counter%",
|
|
format="png",
|
|
quality=100,
|
|
embed_workflow=False,
|
|
save_on_run=True,
|
|
)
|
|
|
|
saved_names = [os.path.basename(path) for path, _args, _kwargs in FakePILImage.saved_paths]
|
|
self.assertEqual(saved_names, ["counter_001.png", "counter_002.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("\\", "/"))
|
|
self.assertTrue(result["ui"]["images"][0]["token"])
|
|
|
|
def test_save_image_handles_windows_different_drive_paths(self):
|
|
node = self.image_nodes.DumasSaveImageNode()
|
|
image = FakeTensorBatch()
|
|
|
|
with mock.patch.object(
|
|
self.image_nodes.folder_paths,
|
|
"get_output_directory",
|
|
return_value="C:\\ComfyUI\\output",
|
|
), mock.patch.object(
|
|
self.image_nodes.os.path,
|
|
"abspath",
|
|
return_value="D:\\renders",
|
|
):
|
|
result = node.save_images(
|
|
images=image,
|
|
folder="D:\\renders",
|
|
pattern="drive_%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"], "D:/renders")
|
|
self.assertTrue(result["ui"]["images"][0]["token"])
|
|
|
|
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()
|