Vito Sansevero
b2ea2597e6
The output scan handed every .mp4 to Pillow and logged one ERROR per file. ComfyUI's video savers store the queued graph as JSON in the container's comment tag, so those files carry prompts worth linking. - utils/video_metadata.py reads the comment/description tag through ffprobe with a timeout, normalised to the PNG text-chunk shape so the existing parser handles both; without ffprobe videos are skipped and a single INFO notice is logged - _extract_comfyui_metadata routes videos there and logs unreadable images as warnings, since the scan carries on past them Test hygiene: every aiohttp test base class overrode tearDownAsync without calling super(), so no test client was ever closed. That produced thousands of ResourceWarnings, "Unclosed client session" errors from asyncio that broke log assertions, and doubled the suite's run time.
400 lines
15 KiB
Python
400 lines
15 KiB
Python
"""
|
|
Tests for GalleryConfig and PromptManagerConfig.
|
|
|
|
The config module imports `from server import PromptServer` at module level,
|
|
which is ComfyUI's server. We mock it before import to run tests standalone.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import MagicMock
|
|
|
|
# Mock ComfyUI's server module before importing config
|
|
_mock_server = MagicMock()
|
|
_mock_server.PromptServer.instance.routes = MagicMock()
|
|
sys.modules.setdefault("server", _mock_server)
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from py.config import GalleryConfig, IntegrationConfig, PromptManagerConfig
|
|
|
|
|
|
class TestGalleryConfig(unittest.TestCase):
|
|
"""Test GalleryConfig class."""
|
|
|
|
def setUp(self):
|
|
"""Save original values to restore after each test."""
|
|
self._orig = GalleryConfig.get_config()
|
|
|
|
def tearDown(self):
|
|
"""Restore original config values."""
|
|
GalleryConfig.update_config(self._orig)
|
|
|
|
def test_get_config_structure(self):
|
|
config = GalleryConfig.get_config()
|
|
self.assertIn("monitoring", config)
|
|
self.assertIn("tracking", config)
|
|
self.assertIn("database", config)
|
|
self.assertIn("web_interface", config)
|
|
self.assertIn("performance", config)
|
|
|
|
def test_default_values(self):
|
|
config = GalleryConfig.get_config()
|
|
self.assertTrue(config["monitoring"]["enabled"])
|
|
self.assertEqual(
|
|
config["monitoring"]["extensions"],
|
|
[".png", ".jpg", ".jpeg", ".webp", ".gif"],
|
|
)
|
|
self.assertEqual(config["tracking"]["prompt_timeout"], 600)
|
|
self.assertEqual(config["web_interface"]["images_per_page"], 20)
|
|
|
|
def test_update_monitoring(self):
|
|
GalleryConfig.update_config(
|
|
{"monitoring": {"enabled": False, "processing_delay": 5.0}}
|
|
)
|
|
self.assertFalse(GalleryConfig.MONITORING_ENABLED)
|
|
self.assertEqual(GalleryConfig.PROCESSING_DELAY, 5.0)
|
|
|
|
def test_update_tracking(self):
|
|
GalleryConfig.update_config({"tracking": {"prompt_timeout": 300}})
|
|
self.assertEqual(GalleryConfig.PROMPT_TIMEOUT, 300)
|
|
|
|
def test_update_database(self):
|
|
GalleryConfig.update_config(
|
|
{"database": {"auto_cleanup": False, "max_image_age_days": 30}}
|
|
)
|
|
self.assertFalse(GalleryConfig.AUTO_CLEANUP_MISSING_FILES)
|
|
self.assertEqual(GalleryConfig.MAX_IMAGE_AGE_DAYS, 30)
|
|
|
|
def test_update_web_interface(self):
|
|
GalleryConfig.update_config(
|
|
{"web_interface": {"images_per_page": 50, "thumbnail_size": 512}}
|
|
)
|
|
self.assertEqual(GalleryConfig.IMAGES_PER_PAGE, 50)
|
|
self.assertEqual(GalleryConfig.THUMBNAIL_SIZE, 512)
|
|
|
|
def test_update_performance(self):
|
|
GalleryConfig.update_config({"performance": {"max_concurrent_processing": 5}})
|
|
self.assertEqual(GalleryConfig.MAX_CONCURRENT_PROCESSING, 5)
|
|
|
|
def test_partial_update_preserves_other_values(self):
|
|
original_timeout = GalleryConfig.PROMPT_TIMEOUT
|
|
GalleryConfig.update_config({"monitoring": {"enabled": False}})
|
|
self.assertEqual(GalleryConfig.PROMPT_TIMEOUT, original_timeout)
|
|
|
|
def test_empty_update_changes_nothing(self):
|
|
before = GalleryConfig.get_config()
|
|
GalleryConfig.update_config({})
|
|
after = GalleryConfig.get_config()
|
|
self.assertEqual(before, after)
|
|
|
|
|
|
class TestPromptManagerConfig(unittest.TestCase):
|
|
"""Test PromptManagerConfig class."""
|
|
|
|
def setUp(self):
|
|
self._orig = PromptManagerConfig.get_config()
|
|
|
|
def tearDown(self):
|
|
PromptManagerConfig.update_config(self._orig)
|
|
|
|
def test_get_config_structure(self):
|
|
config = PromptManagerConfig.get_config()
|
|
self.assertIn("database", config)
|
|
self.assertIn("web_ui", config)
|
|
self.assertIn("performance", config)
|
|
self.assertIn("gallery", config)
|
|
|
|
def test_gallery_config_nested(self):
|
|
config = PromptManagerConfig.get_config()
|
|
# Gallery section should match GalleryConfig output
|
|
self.assertIn("monitoring", config["gallery"])
|
|
self.assertIn("tracking", config["gallery"])
|
|
|
|
def test_default_database_values(self):
|
|
config = PromptManagerConfig.get_config()
|
|
self.assertEqual(config["database"]["default_path"], "prompts.db")
|
|
self.assertTrue(config["database"]["enable_duplicate_detection"])
|
|
self.assertTrue(config["database"]["enable_auto_save"])
|
|
|
|
def test_default_web_ui_values(self):
|
|
config = PromptManagerConfig.get_config()
|
|
self.assertEqual(config["web_ui"]["webui_display_mode"], "newtab")
|
|
|
|
def test_update_database(self):
|
|
PromptManagerConfig.update_config({"database": {"default_path": "custom.db"}})
|
|
self.assertEqual(PromptManagerConfig.DEFAULT_DB_PATH, "custom.db")
|
|
|
|
def test_update_web_ui(self):
|
|
PromptManagerConfig.update_config({"web_ui": {"webui_display_mode": "popup"}})
|
|
self.assertEqual(PromptManagerConfig.WEBUI_DISPLAY_MODE, "popup")
|
|
|
|
def test_update_performance(self):
|
|
PromptManagerConfig.update_config({"performance": {"max_search_results": 50}})
|
|
self.assertEqual(PromptManagerConfig.MAX_SEARCH_RESULTS, 50)
|
|
|
|
def test_worker_threads_default_is_half_the_cores_and_current_is_in_range(self):
|
|
# The current value may come from config.json (loaded at import), so the
|
|
# rule is checked on the constant and only the range on the live value.
|
|
cores = PromptManagerConfig.max_worker_threads()
|
|
self.assertGreaterEqual(cores, 1)
|
|
self.assertEqual(PromptManagerConfig.DEFAULT_WORKER_THREADS, max(1, cores // 2))
|
|
current = PromptManagerConfig.get_config()["performance"]["worker_threads"]
|
|
self.assertGreaterEqual(current, 1)
|
|
self.assertLessEqual(current, cores)
|
|
|
|
def test_update_worker_threads_is_clamped_to_the_machine(self):
|
|
cores = PromptManagerConfig.max_worker_threads()
|
|
PromptManagerConfig.update_config({"performance": {"worker_threads": 1}})
|
|
self.assertEqual(PromptManagerConfig.WORKER_THREADS, 1)
|
|
PromptManagerConfig.update_config(
|
|
{"performance": {"worker_threads": cores + 50}}
|
|
)
|
|
self.assertEqual(PromptManagerConfig.WORKER_THREADS, cores)
|
|
PromptManagerConfig.update_config({"performance": {"worker_threads": 0}})
|
|
self.assertEqual(PromptManagerConfig.WORKER_THREADS, 1)
|
|
|
|
def test_update_worker_threads_ignores_junk_from_a_hand_edited_file(self):
|
|
PromptManagerConfig.update_config({"performance": {"worker_threads": 1}})
|
|
PromptManagerConfig.update_config({"performance": {"worker_threads": "lots"}})
|
|
self.assertEqual(PromptManagerConfig.WORKER_THREADS, 1)
|
|
PromptManagerConfig.update_config({"performance": {"worker_threads": True}})
|
|
self.assertEqual(PromptManagerConfig.WORKER_THREADS, 1)
|
|
|
|
def test_update_propagates_to_gallery(self):
|
|
PromptManagerConfig.update_config(
|
|
{"gallery": {"monitoring": {"enabled": False}}}
|
|
)
|
|
self.assertFalse(GalleryConfig.MONITORING_ENABLED)
|
|
# Restore
|
|
GalleryConfig.MONITORING_ENABLED = True
|
|
|
|
def test_save_and_load_file(self):
|
|
tmp = tempfile.NamedTemporaryFile(
|
|
delete=False, suffix=".json", dir=tempfile.gettempdir()
|
|
)
|
|
tmp.close()
|
|
try:
|
|
# Modify a value
|
|
PromptManagerConfig.MAX_SEARCH_RESULTS = 42
|
|
PromptManagerConfig.save_to_file(tmp.name)
|
|
|
|
# Reset and reload
|
|
PromptManagerConfig.MAX_SEARCH_RESULTS = 100
|
|
PromptManagerConfig.load_from_file(tmp.name)
|
|
self.assertEqual(PromptManagerConfig.MAX_SEARCH_RESULTS, 42)
|
|
finally:
|
|
os.unlink(tmp.name)
|
|
|
|
def test_load_nonexistent_file_uses_defaults(self):
|
|
# Should not raise, just log and continue
|
|
original = PromptManagerConfig.MAX_SEARCH_RESULTS
|
|
PromptManagerConfig.load_from_file("/nonexistent/path/config.json")
|
|
self.assertEqual(PromptManagerConfig.MAX_SEARCH_RESULTS, original)
|
|
|
|
def test_load_invalid_json_file(self):
|
|
tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".json", mode="w")
|
|
tmp.write("not valid json {{")
|
|
tmp.close()
|
|
try:
|
|
original = PromptManagerConfig.MAX_SEARCH_RESULTS
|
|
PromptManagerConfig.load_from_file(tmp.name)
|
|
# Should not crash, keeps existing values
|
|
self.assertEqual(PromptManagerConfig.MAX_SEARCH_RESULTS, original)
|
|
finally:
|
|
os.unlink(tmp.name)
|
|
|
|
def test_save_creates_directory(self):
|
|
tmp_dir = tempfile.mkdtemp()
|
|
nested_path = os.path.join(tmp_dir, "subdir", "config.json")
|
|
try:
|
|
PromptManagerConfig.save_to_file(nested_path)
|
|
self.assertTrue(os.path.exists(nested_path))
|
|
with open(nested_path) as f:
|
|
data = json.load(f)
|
|
self.assertIn("database", data)
|
|
finally:
|
|
import shutil
|
|
|
|
shutil.rmtree(tmp_dir)
|
|
|
|
def test_saved_file_is_private_to_the_user(self):
|
|
import stat
|
|
|
|
tmp_dir = tempfile.mkdtemp()
|
|
try:
|
|
path = os.path.join(tmp_dir, "config.json")
|
|
PromptManagerConfig.save_to_file(path)
|
|
mode = stat.S_IMODE(os.stat(path).st_mode)
|
|
if os.name == "posix":
|
|
self.assertEqual(mode, 0o600)
|
|
with open(path) as f:
|
|
self.assertIn("database", json.load(f))
|
|
# Overwriting an existing, more open file tightens it as well
|
|
os.chmod(path, 0o644)
|
|
PromptManagerConfig.save_to_file(path)
|
|
if os.name == "posix":
|
|
self.assertEqual(stat.S_IMODE(os.stat(path).st_mode), 0o600)
|
|
finally:
|
|
import shutil
|
|
|
|
shutil.rmtree(tmp_dir, ignore_errors=True)
|
|
|
|
def test_saved_file_is_valid_json(self):
|
|
tmp = tempfile.NamedTemporaryFile(
|
|
delete=False, suffix=".json", dir=tempfile.gettempdir()
|
|
)
|
|
tmp.close()
|
|
try:
|
|
PromptManagerConfig.save_to_file(tmp.name)
|
|
with open(tmp.name) as f:
|
|
data = json.load(f)
|
|
self.assertIsInstance(data, dict)
|
|
finally:
|
|
os.unlink(tmp.name)
|
|
|
|
|
|
class TestIntegrationConfig(unittest.TestCase):
|
|
"""Test IntegrationConfig for LoRA Manager settings."""
|
|
|
|
def setUp(self):
|
|
"""Save original values to restore after each test."""
|
|
self._orig = IntegrationConfig.get_config()
|
|
|
|
def tearDown(self):
|
|
"""Restore original config values."""
|
|
IntegrationConfig.update_config(self._orig)
|
|
|
|
def test_reset_to_disabled(self):
|
|
"""Integrations can be fully disabled via update_config."""
|
|
IntegrationConfig.update_config(
|
|
{
|
|
"lora_manager": {
|
|
"enabled": False,
|
|
"path": "",
|
|
"trigger_words_enabled": False,
|
|
"civitai_api_key": "",
|
|
}
|
|
}
|
|
)
|
|
config = IntegrationConfig.get_config()
|
|
lora = config["lora_manager"]
|
|
self.assertFalse(lora["enabled"])
|
|
self.assertEqual(lora["path"], "")
|
|
self.assertFalse(lora["trigger_words_enabled"])
|
|
self.assertEqual(lora["civitai_api_key"], "")
|
|
|
|
def test_get_config_structure(self):
|
|
config = IntegrationConfig.get_config()
|
|
self.assertIn("lora_manager", config)
|
|
lora = config["lora_manager"]
|
|
self.assertIn("enabled", lora)
|
|
self.assertIn("path", lora)
|
|
self.assertIn("trigger_words_enabled", lora)
|
|
self.assertIn("civitai_api_key", lora)
|
|
|
|
def test_update_config_enables(self):
|
|
IntegrationConfig.update_config(
|
|
{
|
|
"lora_manager": {
|
|
"enabled": True,
|
|
"path": "/some/path",
|
|
"trigger_words_enabled": True,
|
|
"civitai_api_key": "test-key-123",
|
|
}
|
|
}
|
|
)
|
|
self.assertTrue(IntegrationConfig.LORA_MANAGER_ENABLED)
|
|
self.assertEqual(IntegrationConfig.LORA_MANAGER_PATH, "/some/path")
|
|
self.assertTrue(IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED)
|
|
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, "test-key-123")
|
|
|
|
def test_update_partial(self):
|
|
"""Updating one field shouldn't affect others."""
|
|
# Reset to known state first
|
|
IntegrationConfig.update_config(
|
|
{"lora_manager": {"enabled": False, "path": "/known"}}
|
|
)
|
|
# Now update only enabled
|
|
IntegrationConfig.update_config({"lora_manager": {"enabled": True}})
|
|
self.assertTrue(IntegrationConfig.LORA_MANAGER_ENABLED)
|
|
self.assertEqual(IntegrationConfig.LORA_MANAGER_PATH, "/known")
|
|
|
|
def test_update_empty_dict_noop(self):
|
|
"""Updating with empty dict preserves current state."""
|
|
before = IntegrationConfig.get_config()
|
|
IntegrationConfig.update_config({})
|
|
after = IntegrationConfig.get_config()
|
|
self.assertEqual(before, after)
|
|
|
|
def test_roundtrip(self):
|
|
"""get_config → update_config → get_config should be stable."""
|
|
IntegrationConfig.update_config(
|
|
{"lora_manager": {"enabled": True, "path": "/test"}}
|
|
)
|
|
config1 = IntegrationConfig.get_config()
|
|
IntegrationConfig.update_config(config1)
|
|
config2 = IntegrationConfig.get_config()
|
|
self.assertEqual(config1, config2)
|
|
|
|
|
|
class TestConfigPathOverride(unittest.TestCase):
|
|
"""PROMPT_MANAGER_CONFIG_PATH relocates the persisted config.json."""
|
|
|
|
ENV = "PROMPT_MANAGER_CONFIG_PATH"
|
|
|
|
def setUp(self):
|
|
self._orig_env = os.environ.pop(self.ENV, None)
|
|
self._orig = PromptManagerConfig.get_config()
|
|
|
|
def tearDown(self):
|
|
if self._orig_env is None:
|
|
os.environ.pop(self.ENV, None)
|
|
else:
|
|
os.environ[self.ENV] = self._orig_env
|
|
PromptManagerConfig.update_config(self._orig)
|
|
|
|
def test_default_config_path_is_repo_config_json(self):
|
|
repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
expected = os.path.join(repo_root, "config.json")
|
|
self.assertEqual(
|
|
os.path.normcase(PromptManagerConfig.get_config_path()),
|
|
os.path.normcase(expected),
|
|
)
|
|
|
|
def test_env_override_is_used(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
custom = os.path.join(tmpdir, "custom.json")
|
|
os.environ[self.ENV] = custom
|
|
self.assertEqual(PromptManagerConfig.get_config_path(), custom)
|
|
|
|
def test_save_without_path_uses_env_override(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
custom = os.path.join(tmpdir, "nested", "custom.json")
|
|
os.environ[self.ENV] = custom
|
|
PromptManagerConfig.RESULT_TIMEOUT = 42
|
|
|
|
PromptManagerConfig.save_to_file()
|
|
|
|
with open(custom) as f:
|
|
saved = json.load(f)
|
|
self.assertEqual(saved["web_ui"]["result_timeout"], 42)
|
|
|
|
def test_load_without_path_uses_env_override(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
custom = os.path.join(tmpdir, "custom.json")
|
|
with open(custom, "w") as f:
|
|
json.dump({"web_ui": {"result_timeout": 77}}, f)
|
|
os.environ[self.ENV] = custom
|
|
|
|
PromptManagerConfig.load_from_file()
|
|
|
|
self.assertEqual(PromptManagerConfig.RESULT_TIMEOUT, 77)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|