diff --git a/prompt_manager_base.py b/prompt_manager_base.py index 75106f8..9a393fc 100644 --- a/prompt_manager_base.py +++ b/prompt_manager_base.py @@ -219,14 +219,10 @@ class PromptManagerBase: } def cleanup_gallery_system(self): - """Clean up gallery system resources.""" - try: - if hasattr(self, "image_monitor"): - self.image_monitor.stop_monitoring() - self.logger.debug("Gallery system cleaned up") - except Exception as e: - self.logger.error(f"Error cleaning up gallery system: {e}") + """Clean up gallery system resources. - def __del__(self): - """Cleanup when object is destroyed.""" - self.cleanup_gallery_system() + Note: The image monitor is a singleton shared across all node instances, + so we intentionally do NOT stop it here. It should run for the entire + ComfyUI session. Stopping it would kill monitoring for all nodes. + """ + self.logger.debug("Gallery system cleaned up (monitor continues running)") diff --git a/py/api/admin.py b/py/api/admin.py index b4060c6..072deb7 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -497,7 +497,7 @@ class AdminRoutesMixin: # Check database try: - db_path = "prompts.db" + db_path = self.db.model.db_path if os.path.exists(db_path): with sqlite3.connect(db_path) as conn: conn.row_factory = sqlite3.Row @@ -559,12 +559,38 @@ class AdminRoutesMixin: # Check output directories output_dirs = [] - potential_dirs = ["output", "../output", "../../output"] - for dir_path in potential_dirs: - abs_path = os.path.abspath(dir_path) - if os.path.exists(abs_path): - output_dirs.append(abs_path) + # Check user-configured gallery directories from config.json + try: + from ..config import GalleryConfig + + for cfg_dir in GalleryConfig.MONITORING_DIRECTORIES: + abs_path = os.path.abspath(cfg_dir) + if os.path.exists(abs_path) and abs_path not in output_dirs: + output_dirs.append(abs_path) + except Exception as e: + self.logger.debug(f"Could not load gallery config: {e}") + + # Check ComfyUI's own output directory (respects --output-directory) + try: + import folder_paths + + comfyui_output = folder_paths.get_output_directory() + if comfyui_output and os.path.exists(comfyui_output): + abs_path = os.path.abspath(comfyui_output) + if abs_path not in output_dirs: + output_dirs.append(abs_path) + except ImportError: + self.logger.debug( + "folder_paths not available, skipping ComfyUI output dir" + ) + + # Fallback: check common relative paths + if not output_dirs: + for dir_path in ["output", "../output", "../../output"]: + abs_path = os.path.abspath(dir_path) + if os.path.exists(abs_path) and abs_path not in output_dirs: + output_dirs.append(abs_path) results["comfyui_output"] = { "status": "ok" if output_dirs else "warning", @@ -801,7 +827,7 @@ class AdminRoutesMixin: async def backup_database(self, request): """Backup the entire prompts.db database file.""" try: - db_path = "prompts.db" + db_path = self.db.model.db_path if not os.path.exists(db_path): return web.json_response( @@ -894,7 +920,7 @@ class AdminRoutesMixin: cursor = conn.execute("SELECT COUNT(*) as count FROM prompts") prompt_count = cursor.fetchone()["count"] - db_path = "prompts.db" + db_path = self.db.model.db_path backup_path = f"{db_path}.backup_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}" if os.path.exists(db_path): diff --git a/tests/test_diagnostics_config.py b/tests/test_diagnostics_config.py new file mode 100644 index 0000000..cc69437 --- /dev/null +++ b/tests/test_diagnostics_config.py @@ -0,0 +1,189 @@ +""" +Tests for issue #132: Diagnostics, backup, and restore must respect config.json. + +Verifies that: +- GalleryDiagnostics resolves db_path from config (not hardcoded "prompts.db") +- The admin API diagnostics/backup/restore endpoints use the configured db_path +- Output directory checks consult GalleryConfig before falling back to heuristics +""" + +import asyncio +import json +import os +import sys +import tempfile +import types +import unittest +from unittest.mock import patch, MagicMock + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +EXTENSION_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +# --------------------------------------------------------------------------- +# GalleryDiagnostics path resolution +# --------------------------------------------------------------------------- +class TestGalleryDiagnosticsDbPath(unittest.TestCase): + """GalleryDiagnostics should resolve db_path from config, not hardcode it.""" + + def test_default_resolves_from_config(self): + """GalleryDiagnostics() with no args should read config like PromptDatabase.""" + from utils.diagnostics import GalleryDiagnostics + + with tempfile.TemporaryDirectory() as tmpdir: + custom_path = os.path.join(tmpdir, "custom.db") + fake_config_mod = types.ModuleType("py.config") + fake_pm_config = type( + "PromptManagerConfig", (), {"DEFAULT_DB_PATH": custom_path} + ) + fake_config_mod.PromptManagerConfig = fake_pm_config + + with patch.dict(sys.modules, {"py.config": fake_config_mod}): + diag = GalleryDiagnostics() + + self.assertEqual(diag.db_path, custom_path) + + def test_default_without_config_falls_back_to_extension_root(self): + """Without config, should default to prompts.db in extension root.""" + from utils.diagnostics import GalleryDiagnostics + + diag = GalleryDiagnostics() + expected = os.path.join(EXTENSION_ROOT, "prompts.db") + self.assertEqual(diag.db_path, expected) + + def test_explicit_path_overrides_config(self): + """An explicit db_path argument should be used as-is.""" + from utils.diagnostics import GalleryDiagnostics + + diag = GalleryDiagnostics("/explicit/path.db") + self.assertEqual(diag.db_path, "/explicit/path.db") + + def test_config_relative_path_resolved_to_extension_root(self): + """A relative path from config should resolve against extension root.""" + from utils.diagnostics import GalleryDiagnostics + + fake_config_mod = types.ModuleType("py.config") + fake_pm_config = type( + "PromptManagerConfig", (), {"DEFAULT_DB_PATH": "data/prompts.db"} + ) + fake_config_mod.PromptManagerConfig = fake_pm_config + + with patch.dict(sys.modules, {"py.config": fake_config_mod}): + diag = GalleryDiagnostics() + + expected = os.path.join(EXTENSION_ROOT, "data", "prompts.db") + self.assertEqual(diag.db_path, expected) + + +# --------------------------------------------------------------------------- +# Admin API endpoints use self.db.model.db_path +# --------------------------------------------------------------------------- +class TestAdminEndpointsUseConfigPath(unittest.TestCase): + """Admin diagnostics, backup, and restore must use the configured db_path.""" + + def _make_api_stub(self, db_path): + """Create a minimal object that mimics the admin mixin's self.""" + stub = MagicMock() + stub.db.model.db_path = db_path + stub.logger = MagicMock() + return stub + + def _run_async(self, coro): + """Run a coroutine on a fresh event loop to avoid aiohttp pollution.""" + loop = asyncio.new_event_loop() + try: + return loop.run_until_complete(coro) + finally: + loop.close() + + def test_run_diagnostics_uses_model_db_path(self): + """run_diagnostics should check self.db.model.db_path, not 'prompts.db'.""" + if "server" not in sys.modules: + sys.modules["server"] = MagicMock() + + from py.api.admin import AdminRoutesMixin + + with tempfile.NamedTemporaryFile(delete=False, suffix=".db") as f: + db_file = f.name + + try: + import sqlite3 + + with sqlite3.connect(db_file) as conn: + conn.execute( + "CREATE TABLE prompts (id INTEGER PRIMARY KEY, text TEXT, created_at TEXT)" + ) + conn.execute( + "INSERT INTO prompts (text, created_at) VALUES ('test', '2024-01-01')" + ) + + stub = self._make_api_stub(db_file) + result = self._run_async( + AdminRoutesMixin.run_diagnostics(stub, MagicMock()) + ) + + body = json.loads(result.body) + self.assertTrue(body["success"]) + self.assertEqual(body["diagnostics"]["database"]["status"], "ok") + self.assertEqual(body["diagnostics"]["database"]["prompt_count"], 1) + finally: + os.unlink(db_file) + + def test_run_diagnostics_reports_missing_custom_path(self): + """If the configured db_path doesn't exist, diagnostics should report error.""" + if "server" not in sys.modules: + sys.modules["server"] = MagicMock() + + from py.api.admin import AdminRoutesMixin + + stub = self._make_api_stub("/nonexistent/custom/prompts.db") + + result = self._run_async(AdminRoutesMixin.run_diagnostics(stub, MagicMock())) + body = json.loads(result.body) + self.assertTrue(body["success"]) + self.assertEqual(body["diagnostics"]["database"]["status"], "error") + self.assertIn( + "/nonexistent/custom/prompts.db", + body["diagnostics"]["database"]["message"], + ) + + def test_backup_uses_model_db_path(self): + """backup_database should read from self.db.model.db_path.""" + if "server" not in sys.modules: + sys.modules["server"] = MagicMock() + + from py.api.admin import AdminRoutesMixin + + with tempfile.NamedTemporaryFile(delete=False, suffix=".db") as f: + db_file = f.name + f.write(b"SQLite format 3\x00test data for backup") + + try: + stub = self._make_api_stub(db_file) + result = self._run_async( + AdminRoutesMixin.backup_database(stub, MagicMock()) + ) + + self.assertEqual(result.content_type, "application/octet-stream") + self.assertGreater(len(result.body), 0) + finally: + os.unlink(db_file) + + def test_backup_reports_missing_custom_path(self): + """backup_database should return 404 when configured path doesn't exist.""" + if "server" not in sys.modules: + sys.modules["server"] = MagicMock() + + from py.api.admin import AdminRoutesMixin + + stub = self._make_api_stub("/nonexistent/custom/prompts.db") + + result = self._run_async(AdminRoutesMixin.backup_database(stub, MagicMock())) + body = json.loads(result.body) + self.assertFalse(body["success"]) + self.assertEqual(result.status, 404) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_image_monitor_singleton.py b/tests/test_image_monitor_singleton.py new file mode 100644 index 0000000..08d6fed --- /dev/null +++ b/tests/test_image_monitor_singleton.py @@ -0,0 +1,164 @@ +""" +Tests for image monitor singleton lifecycle. + +Verifies that: +- The singleton image monitor survives node garbage collection +- Node __del__ / cleanup_gallery_system does NOT stop the shared monitor +- Multiple node instances share the same monitor +- The monitor observer stays alive across node lifecycles +""" + +import os +import sys +import threading +import unittest +from unittest.mock import MagicMock, patch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +# Mock ComfyUI server before importing anything that touches config +_mock_server = MagicMock() +_mock_server.PromptServer.instance.routes = MagicMock() +sys.modules["server"] = _mock_server + +# Mock comfyui_integration to avoid import issues +sys.modules["utils.comfyui_integration"] = MagicMock() + +from utils.image_monitor import get_image_monitor +import utils.image_monitor as im_mod + + +class TestMonitorSingletonLifecycle(unittest.TestCase): + """The singleton monitor must survive node instance garbage collection.""" + + def setUp(self): + # Reset singleton for each test + im_mod._monitor_instance = None + + def tearDown(self): + im_mod._monitor_instance = None + + def test_get_image_monitor_returns_singleton(self): + """Multiple calls should return the exact same instance.""" + db = MagicMock() + tracker = MagicMock() + + m1 = get_image_monitor(db, tracker) + m2 = get_image_monitor(db, tracker) + + self.assertIs(m1, m2) + + def test_singleton_survives_across_different_callers(self): + """Different db/tracker args on subsequent calls still return the same instance.""" + db1, tracker1 = MagicMock(), MagicMock() + db2, tracker2 = MagicMock(), MagicMock() + + m1 = get_image_monitor(db1, tracker1) + m2 = get_image_monitor(db2, tracker2) + + self.assertIs(m1, m2) + + def test_cleanup_does_not_stop_monitor(self): + """PromptManagerBase.cleanup_gallery_system must NOT stop the monitor.""" + from prompt_manager_base import PromptManagerBase + + with patch.object(PromptManagerBase, "__init__", lambda self, **kw: None): + node = PromptManagerBase() + node.logger = MagicMock() + + mock_monitor = MagicMock() + node.image_monitor = mock_monitor + + node.cleanup_gallery_system() + + mock_monitor.stop_monitoring.assert_not_called() + + def test_del_does_not_stop_monitor(self): + """Node __del__ must NOT stop the singleton monitor.""" + from prompt_manager_base import PromptManagerBase + + with patch.object(PromptManagerBase, "__init__", lambda self, **kw: None): + node = PromptManagerBase() + node.logger = MagicMock() + + mock_monitor = MagicMock() + node.image_monitor = mock_monitor + + # Simulate garbage collection + del node + + mock_monitor.stop_monitoring.assert_not_called() + + def test_observer_stays_alive_after_node_cleanup(self): + """A running observer must remain alive after node cleanup.""" + db = MagicMock() + tracker = MagicMock() + monitor = get_image_monitor(db, tracker) + + # Simulate a running observer + mock_observer = MagicMock() + mock_observer.is_alive.return_value = True + monitor.observer = mock_observer + monitor.handler = MagicMock() + monitor.monitored_directories = ["/fake/output"] + + from prompt_manager_base import PromptManagerBase + + with patch.object(PromptManagerBase, "__init__", lambda self, **kw: None): + node = PromptManagerBase() + node.logger = MagicMock() + node.image_monitor = monitor + + node.cleanup_gallery_system() + + # Observer must still be alive + self.assertIsNotNone(monitor.observer) + self.assertTrue(monitor.observer.is_alive()) + self.assertEqual(monitor.monitored_directories, ["/fake/output"]) + + def test_monitor_start_not_called_when_already_running(self): + """start_monitoring should be a no-op if observer is already active.""" + db = MagicMock() + tracker = MagicMock() + monitor = get_image_monitor(db, tracker) + + # Set up as if already running + mock_observer = MagicMock() + monitor.observer = mock_observer + + monitor.start_monitoring() + + # Should not create a new observer + self.assertIs(monitor.observer, mock_observer) + + +class TestMonitorThreadSafety(unittest.TestCase): + """Singleton creation must be thread-safe.""" + + def setUp(self): + im_mod._monitor_instance = None + + def tearDown(self): + im_mod._monitor_instance = None + + def test_concurrent_get_image_monitor_returns_same_instance(self): + """Multiple threads calling get_image_monitor must get the same instance.""" + results = [] + barrier = threading.Barrier(5) + + def get_monitor(): + barrier.wait() + m = get_image_monitor(MagicMock(), MagicMock()) + results.append(id(m)) + + threads = [threading.Thread(target=get_monitor) for _ in range(5)] + for t in threads: + t.start() + for t in threads: + t.join() + + self.assertEqual(len(set(results)), 1, "All threads must get the same instance") + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/diagnostics.py b/utils/diagnostics.py index 60bb7c6..ffed61a 100644 --- a/utils/diagnostics.py +++ b/utils/diagnostics.py @@ -58,15 +58,34 @@ class GalleryDiagnostics: - Additional data specific to the diagnostic """ - def __init__(self, db_path: str = "prompts.db"): + def __init__(self, db_path: str = None): """Initialize the diagnostics system. Args: - db_path: Path to the SQLite database file to diagnose + db_path: Path to the SQLite database file to diagnose. + If None, resolves from config.json or defaults to + prompts.db in the extension root. """ + if db_path is None: + db_path = self._resolve_db_path() self.db_path = db_path self.logger = get_logger("prompt_manager.diagnostics") + @staticmethod + def _resolve_db_path() -> str: + """Resolve database path from config, matching PromptDatabase logic.""" + extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + try: + from py.config import PromptManagerConfig + + db_path = PromptManagerConfig.DEFAULT_DB_PATH + except Exception: + db_path = "prompts.db" + + if not os.path.isabs(db_path): + db_path = os.path.join(extension_root, db_path) + return db_path + def run_full_diagnostic(self) -> Dict[str, Any]: """Run a complete diagnostic check. @@ -268,23 +287,10 @@ class GalleryDiagnostics: output_dirs = [] - # Try to detect ComfyUI output directories - potential_dirs = [ - "output", - "../output", - "../../output", - "ComfyUI/output", - "../ComfyUI/output", - "../../ComfyUI/output", - ] - - for dir_path in potential_dirs: - abs_path = os.path.abspath(dir_path) - if os.path.exists(abs_path): + def _add_dir(abs_path, source): + if abs_path not in output_dirs: output_dirs.append(abs_path) - self.logger.info(f" [DIR] Found output dir: {abs_path}") - - # Count images in this directory + self.logger.info(f" [DIR] {source}: {abs_path}") try: image_files = [] for ext in [".png", ".jpg", ".jpeg", ".webp"]: @@ -293,17 +299,41 @@ class GalleryDiagnostics: except Exception as e: self.logger.error(f" [FAIL] Error scanning: {e}") - # Try ComfyUI's folder_paths + # Check user-configured gallery directories from config.json + try: + from py.config import GalleryConfig + + for cfg_dir in GalleryConfig.MONITORING_DIRECTORIES: + abs_path = os.path.abspath(cfg_dir) + if os.path.exists(abs_path): + _add_dir(abs_path, "Configured directory") + except Exception as e: + self.logger.debug(f" [NOTE] Could not load gallery config: {e}") + + # Check ComfyUI's own output directory (respects --output-directory) try: import folder_paths comfyui_output = folder_paths.get_output_directory() - if comfyui_output and comfyui_output not in output_dirs: - output_dirs.append(comfyui_output) - self.logger.info(f" [DIR] ComfyUI output dir: {comfyui_output}") + if comfyui_output and os.path.exists(comfyui_output): + _add_dir(os.path.abspath(comfyui_output), "ComfyUI output dir") except ImportError: self.logger.warning(" [WARN] ComfyUI folder_paths not available") + # Fallback: check common relative paths + if not output_dirs: + for dir_path in [ + "output", + "../output", + "../../output", + "ComfyUI/output", + "../ComfyUI/output", + "../../ComfyUI/output", + ]: + abs_path = os.path.abspath(dir_path) + if os.path.exists(abs_path): + _add_dir(abs_path, "Found output dir") + return { "status": "ok" if output_dirs else "warning", "message": "No output directories found" if not output_dirs else None,