163 lines
6.7 KiB
Python
163 lines
6.7 KiB
Python
import os
|
|
import shutil
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from typing import cast
|
|
from unittest.mock import ANY, MagicMock, patch
|
|
|
|
import torch
|
|
|
|
# Mock ComfyUI dependencies
|
|
sys.modules["folder_paths"] = MagicMock()
|
|
sys.modules["server"] = MagicMock()
|
|
sys.modules["comfy.cli_args"] = MagicMock()
|
|
sys.modules["nodes"] = MagicMock()
|
|
sys.modules["comfy"] = MagicMock()
|
|
sys.modules["comfy.sd"] = MagicMock()
|
|
sys.modules["comfy.utils"] = MagicMock()
|
|
sys.modules["comfy.samplers"] = MagicMock()
|
|
|
|
from comfy.cli_args import args # noqa: E402
|
|
|
|
args.disable_metadata = False
|
|
|
|
# Import test target
|
|
import py.image_manager.common.db.schema as database # noqa: E402
|
|
import py.meld_save_image.nodes as meld_node # noqa: E402
|
|
|
|
|
|
class TestImageManager(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
# Create a temporary directory for DB and output
|
|
self.test_dir = tempfile.mkdtemp()
|
|
self.db_path = os.path.join(self.test_dir, "test.db")
|
|
|
|
# Patch DB_PATH in database module and db.client module
|
|
self.patcher_db = patch("py.image_manager.common.db.client.DB_PATH", self.db_path)
|
|
self.patcher_db.start()
|
|
self.patcher_trash = patch("py.image_manager.common.db.client.TRASH_DIR", os.path.join(self.test_dir, "trash"))
|
|
self.patcher_trash.start()
|
|
|
|
# Initialize DB
|
|
database.init_db()
|
|
|
|
# Mock folder_paths in nodes module
|
|
self.mock_output_dir = os.path.join(self.test_dir, "output")
|
|
os.makedirs(self.mock_output_dir, exist_ok=True)
|
|
meld_node.folder_paths.get_output_directory.return_value = self.mock_output_dir # type: ignore
|
|
|
|
# Reset get_save_image_path mock
|
|
meld_node.folder_paths.get_save_image_path.reset_mock(side_effect=True, return_value=True) # type: ignore
|
|
|
|
# Instance of node
|
|
self.node = meld_node.MeldSaveImage()
|
|
|
|
def tearDown(self) -> None:
|
|
self.patcher_db.stop()
|
|
self.patcher_trash.stop()
|
|
shutil.rmtree(self.test_dir)
|
|
|
|
def test_save_images_filename_prefix_passed(self) -> None:
|
|
"""Verify that filename_prefix is correctly passed to get_save_image_path"""
|
|
images = torch.zeros((1, 64, 64, 3))
|
|
mock_get_save_path = cast(MagicMock, meld_node.folder_paths.get_save_image_path) # type: ignore
|
|
mock_get_save_path.return_value = (self.mock_output_dir, "Meld", 1, "", "Meld")
|
|
|
|
self.node.save_images(images=images, filename_prefix="CustomPrefix")
|
|
|
|
# Check if get_save_image_path was called with the correct prefix
|
|
mock_get_save_path.assert_called_with("CustomPrefix", self.mock_output_dir, 64, 64)
|
|
|
|
def test_save_images_db_registration(self) -> None:
|
|
"""Verify that image is registered in DB with correct filename"""
|
|
images = torch.zeros((1, 64, 64, 3))
|
|
mock_get_save_path = cast(MagicMock, meld_node.folder_paths.get_save_image_path) # type: ignore
|
|
mock_get_save_path.return_value = (self.mock_output_dir, "TestFile", 5, "sub", "Test")
|
|
|
|
with patch("PIL.Image.Image.save"):
|
|
self.node.save_images(images=images, filename_prefix="Test")
|
|
|
|
conn = sqlite3.connect(self.db_path)
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT filename, subfolder FROM images")
|
|
row = cursor.fetchone()
|
|
conn.close()
|
|
|
|
# f"{filename_with_batch_num}_{counter:05}_.png"
|
|
# TestFile_00005_.png
|
|
self.assertEqual(row[0], "TestFile_00005_.png")
|
|
self.assertEqual(row[1], "sub")
|
|
|
|
def test_filename_prefix_with_batch_num_token(self) -> None:
|
|
"""Verify that %batch_num% in the filename returned by get_save_image_path is resolved"""
|
|
images = torch.zeros((2, 64, 64, 3)) # Batch of 2
|
|
mock_get_save_path = cast(MagicMock, meld_node.folder_paths.get_save_image_path) # type: ignore
|
|
mock_get_save_path.return_value = (self.mock_output_dir, "Batch_%batch_num%", 1, "", "Batch")
|
|
|
|
with patch("PIL.Image.Image.save"):
|
|
self.node.save_images(images=images, filename_prefix="Batch")
|
|
|
|
conn = sqlite3.connect(self.db_path)
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT filename FROM images ORDER BY id ASC")
|
|
rows = cursor.fetchall()
|
|
conn.close()
|
|
|
|
self.assertEqual(len(rows), 2)
|
|
self.assertEqual(rows[0][0], "Batch_0_00001_.png")
|
|
self.assertEqual(rows[1][0], "Batch_1_00002_.png")
|
|
|
|
def test_save_images_format_fix(self) -> None:
|
|
"""Verify that the trailing underscore is included in the filename for ComfyUI compatibility"""
|
|
images = torch.zeros((1, 64, 64, 3))
|
|
mock_get_save_path = cast(MagicMock, meld_node.folder_paths.get_save_image_path) # type: ignore
|
|
mock_get_save_path.return_value = (self.mock_output_dir, "TestFile", 1, "", "Test")
|
|
|
|
with patch("PIL.Image.Image.save"):
|
|
self.node.save_images(images=images, filename_prefix="Test")
|
|
|
|
conn = sqlite3.connect(self.db_path)
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT filename FROM images")
|
|
row = cursor.fetchone()
|
|
conn.close()
|
|
|
|
# Should be TestFile_00001_.png for ComfyUI counter logic to work
|
|
self.assertEqual(row[0], "TestFile_00001_.png")
|
|
|
|
def test_filename_prefix_date_resolution(self) -> None:
|
|
"""Verify that %date:format% and %date% tokens are correctly resolved"""
|
|
images = torch.zeros((1, 64, 64, 3))
|
|
mock_get_save_path = cast(MagicMock, meld_node.folder_paths.get_save_image_path) # type: ignore
|
|
mock_get_save_path.return_value = (self.mock_output_dir, "resolved", 1, "", "resolved")
|
|
|
|
from datetime import datetime
|
|
|
|
expected_date = datetime.now().strftime("%Y-%m-%d")
|
|
|
|
# Test %date%
|
|
self.node.save_images(images=images, filename_prefix="test_%date%")
|
|
mock_get_save_path.assert_called_with(f"test_{expected_date}", ANY, ANY, ANY)
|
|
|
|
# Test %date:yyyy-MM-dd%
|
|
mock_get_save_path.reset_mock()
|
|
self.node.save_images(images=images, filename_prefix="test_%date:yyyy-MM-dd%")
|
|
mock_get_save_path.assert_called_with(f"test_{expected_date}", ANY, ANY, ANY)
|
|
|
|
def test_save_images_returns_images(self) -> None:
|
|
"""Verify that save_images returns the input images tensor"""
|
|
images = torch.zeros((1, 64, 64, 3))
|
|
mock_get_save_path = cast(MagicMock, meld_node.folder_paths.get_save_image_path) # type: ignore
|
|
mock_get_save_path.return_value = (self.mock_output_dir, "test", 1, "", "test")
|
|
|
|
with patch("PIL.Image.Image.save"):
|
|
result = self.node.save_images(images=images)
|
|
self.assertIn("result", result)
|
|
self.assertTrue(torch.equal(result["result"][0], images))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|