Vito Sansevero
489d449ec8
- py/model_gate.py: one process-wide, non-blocking gate for loading,
swapping, running, downloading and unloading a tagging model. Worker
calls run shielded; a release requested while one still runs takes
effect only when it exits (client disconnect keeps the gate).
- Legacy /autotag/single, /start, /download/{type} and /unload hold the
gate and answer 409 {code: "model-busy"} while another operation holds
it; their bodies are otherwise unchanged. /autotag/models gains
suggestion_limit.
- AutoTagService.generate_tag_suggestions: WD14 (name, probability) best
first, GGUF/HF (name, None) in output order, first `limit` distinct
case-insensitively; generate_tags keeps its return; the image file is
closed before inference returns.
- POST /prompt_manager/image-annotations/suggest runs one engine on the
pinned capture, with strict per-engine options, never downloads
(409 model-unavailable) and never writes.
- config.json autotag.defaults.suggestion_limit, integer 1..100, default
10, validated on load.
- Capture probe decides Generate from the image header alone, so prompt
metadata over the inspection budgets no longer blocks it.
478 lines
19 KiB
Python
478 lines
19 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 import config as config_module
|
|
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_infinite_scroll_default_off_and_only_booleans_apply(self):
|
|
self.assertIn("infinite_scroll", PromptManagerConfig.get_config()["web_ui"])
|
|
PromptManagerConfig.update_config({"web_ui": {"infinite_scroll": True}})
|
|
self.assertTrue(PromptManagerConfig.INFINITE_SCROLL)
|
|
PromptManagerConfig.update_config({"web_ui": {"infinite_scroll": "yes"}})
|
|
self.assertTrue(PromptManagerConfig.INFINITE_SCROLL)
|
|
PromptManagerConfig.update_config({"web_ui": {"infinite_scroll": False}})
|
|
self.assertFalse(PromptManagerConfig.INFINITE_SCROLL)
|
|
|
|
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)
|
|
|
|
|
|
class TestAutotagConfig(unittest.TestCase):
|
|
"""``autotag.defaults.suggestion_limit``: integer 1..100, default 10 (D68)."""
|
|
|
|
def setUp(self):
|
|
self.config_class = getattr(config_module, "AutotagConfig", None)
|
|
self.assertIsNotNone(self.config_class, "AutotagConfig is not implemented")
|
|
self._orig = PromptManagerConfig.get_config()
|
|
|
|
def tearDown(self):
|
|
self.config_class.SUGGESTION_LIMIT = 10
|
|
PromptManagerConfig.update_config(self._orig)
|
|
|
|
def _load(self, value):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
path = os.path.join(tmpdir, "config.json")
|
|
with open(path, "w") as f:
|
|
# json writes NaN for float("nan"), which json.load reads back.
|
|
json.dump({"autotag": {"defaults": {"suggestion_limit": value}}}, f)
|
|
with self.assertLogs("prompt_manager.config", level="WARNING") as logs:
|
|
PromptManagerConfig.load_from_file(path)
|
|
return logs
|
|
|
|
def test_default_is_ten_and_persisted_under_autotag_defaults(self):
|
|
self.assertEqual(config_module.DEFAULT_SUGGESTION_LIMIT, 10)
|
|
self.assertEqual(config_module.MAX_SUGGESTION_LIMIT, 100)
|
|
self.assertEqual(self.config_class.suggestion_limit(), 10)
|
|
self.assertEqual(
|
|
PromptManagerConfig.get_config()["autotag"],
|
|
{"defaults": {"suggestion_limit": 10}},
|
|
)
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
path = os.path.join(tmpdir, "config.json")
|
|
PromptManagerConfig.update_config(
|
|
{"autotag": {"defaults": {"suggestion_limit": 25}}}
|
|
)
|
|
PromptManagerConfig.save_to_file(path)
|
|
self.config_class.SUGGESTION_LIMIT = 10
|
|
PromptManagerConfig.load_from_file(path)
|
|
self.assertEqual(self.config_class.suggestion_limit(), 25)
|
|
|
|
def test_valid_limits_are_kept(self):
|
|
for value in (1, 10, 100):
|
|
with self.subTest(value=value):
|
|
PromptManagerConfig.update_config(
|
|
{"autotag": {"defaults": {"suggestion_limit": value}}}
|
|
)
|
|
self.assertEqual(self.config_class.suggestion_limit(), value)
|
|
|
|
def test_suggestion_limit_setting_validated_on_load(self):
|
|
for value in (0, 101, "10", float("nan"), 2.5, True, None, [10]):
|
|
with self.subTest(value=value):
|
|
self.config_class.SUGGESTION_LIMIT = 42
|
|
logs = self._load(value)
|
|
self.assertEqual(self.config_class.suggestion_limit(), 10)
|
|
self.assertIn("suggestion_limit", "\n".join(logs.output))
|
|
|
|
def test_missing_or_malformed_sections_keep_the_current_value(self):
|
|
self.config_class.SUGGESTION_LIMIT = 7
|
|
for config in ({}, {"autotag": {}}, {"autotag": {"defaults": {}}}):
|
|
PromptManagerConfig.update_config(config)
|
|
self.assertEqual(self.config_class.suggestion_limit(), 7)
|
|
with self.assertLogs("prompt_manager.config", level="WARNING"):
|
|
PromptManagerConfig.update_config({"autotag": "nope"})
|
|
with self.assertLogs("prompt_manager.config", level="WARNING"):
|
|
PromptManagerConfig.update_config({"autotag": {"defaults": [1]}})
|
|
self.assertEqual(self.config_class.suggestion_limit(), 7)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|