diff --git a/py/api/__init__.py b/py/api/__init__.py index b3f135d..f243b6b 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -264,12 +264,12 @@ class PromptManagerAPI( db (PromptDatabase): Database connection and operations handler """ - def __init__(self): - """Initialize the PromptManager API with database connection and cleanup.""" + def __init__(self, db: PromptDatabase | None = None) -> None: + """Initialize the API, using an injected database before any cleanup.""" self.logger = get_logger("prompt_manager.api") self.logger.info("Initializing PromptManager API") - self.db = PromptDatabase() + self.db = db if db is not None else PromptDatabase() self._cached_output_dir = None # Lazy-cached by _find_comfyui_output_dir() self._html_cache = {} # Cached HTML file contents keyed by path self._gallery_cache = {} # dict: path_str -> (files, timestamp) diff --git a/tests/test_api.py b/tests/test_api.py index f093393..051881a 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -31,8 +31,7 @@ class APITestCase(AioHTTPTestCase): app = web.Application() routes = web.RouteTableDef() - self.api = PromptManagerAPI() - self.api.db = PromptDatabase(self._temp_db.name) + self.api = PromptManagerAPI(db=PromptDatabase(self._temp_db.name)) self.api.add_routes(routes) app.router.add_routes(routes) return app diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index 870e093..1f91093 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -70,8 +70,7 @@ class AdminAPITestCase(AioHTTPTestCase): db_path = os.path.join(self.tmpdir, "prompts.db") app = web.Application() routes = web.RouteTableDef() - self.api = PromptManagerAPI() - self.api.db = PromptDatabase(db_path) + self.api = PromptManagerAPI(db=PromptDatabase(db_path)) self.api.add_routes(routes) app.router.add_routes(routes) return app diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index 8db42d9..832221d 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -119,8 +119,7 @@ class AutotagAPITestCase(AioHTTPTestCase): app = web.Application() routes = web.RouteTableDef() - self.api = PromptManagerAPI() - self.api.db = PromptDatabase(self._temp_db.name) + self.api = PromptManagerAPI(db=PromptDatabase(self._temp_db.name)) self.api._get_all_output_dirs = lambda: [self.output_dir] self.api.add_routes(routes) app.router.add_routes(routes) diff --git a/tests/test_api_backup_restore.py b/tests/test_api_backup_restore.py index fb473d9..eb55b47 100644 --- a/tests/test_api_backup_restore.py +++ b/tests/test_api_backup_restore.py @@ -45,8 +45,7 @@ class BackupRestoreTestCase(AioHTTPTestCase): app = web.Application() routes = web.RouteTableDef() - self.api = PromptManagerAPI() - self.api.db = PromptDatabase(self.db_path) + self.api = PromptManagerAPI(db=PromptDatabase(self.db_path)) self.api.add_routes(routes) app.router.add_routes(routes) return app diff --git a/tests/test_api_core.py b/tests/test_api_core.py index 16d4a55..91c4bce 100644 --- a/tests/test_api_core.py +++ b/tests/test_api_core.py @@ -322,7 +322,9 @@ class TestLongJobLock(unittest.TestCase): """_acquire_long_job: one job per family at a time, released on exit.""" def setUp(self): - self.api = PromptManagerAPI() + database = MagicMock() + database.cleanup_duplicates.return_value = 0 + self.api = PromptManagerAPI(db=database) def _run(self, coro): import asyncio @@ -373,8 +375,9 @@ class CoreRoutesTestCase(AioHTTPTestCase): self.tmpdir = tempfile.mkdtemp() app = web.Application() routes = web.RouteTableDef() - self.api = PromptManagerAPI() - self.api.db = PromptDatabase(os.path.join(self.tmpdir, "prompts.db")) + self.api = PromptManagerAPI( + db=PromptDatabase(os.path.join(self.tmpdir, "prompts.db")) + ) self.api.add_routes(routes) app.router.add_routes(routes) return app @@ -565,8 +568,9 @@ class TestInstanceHelpers(FolderPathsFixture): def setUp(self): super().setUp() - self.api = PromptManagerAPI() - self.api.db = PromptDatabase(os.path.join(self.tmpdir, "prompts.db")) + self.api = PromptManagerAPI( + db=PromptDatabase(os.path.join(self.tmpdir, "prompts.db")) + ) self.addCleanup(self.api.db.close_all) # runs before the rmtree cleanup orig_dirs = list(GalleryConfig.MONITORING_DIRECTORIES) self.addCleanup(setattr, GalleryConfig, "MONITORING_DIRECTORIES", orig_dirs) @@ -668,7 +672,9 @@ class TestInstanceHelpers(FolderPathsFixture): class TestMetadataParsing(unittest.TestCase): def setUp(self): - self.api = PromptManagerAPI() + database = MagicMock() + database.cleanup_duplicates.return_value = 0 + self.api = PromptManagerAPI(db=database) self.tmpdir = tempfile.mkdtemp() self.addCleanup(shutil.rmtree, self.tmpdir, True) diff --git a/tests/test_api_database_isolation.py b/tests/test_api_database_isolation.py new file mode 100644 index 0000000..75593f4 --- /dev/null +++ b/tests/test_api_database_isolation.py @@ -0,0 +1,71 @@ +"""API fixtures never open the configured database before installing test state.""" + +import unittest.mock + +from py.api import PromptManagerAPI + + +class TestAPIDatabaseIsolation(unittest.TestCase): + def test_injected_database_never_constructs_the_configured_database(self): + database = unittest.mock.MagicMock() + database.cleanup_duplicates.return_value = 0 + with ( + unittest.mock.patch( + "py.api.PromptDatabase", side_effect=AssertionError("default DB opened") + ) as factory, + unittest.mock.patch("py.api.get_logger"), + ): + api = PromptManagerAPI(db=database) + self.assertIs(api.db, database) + factory.assert_not_called() + database.cleanup_duplicates.assert_called_once_with() + + def test_falsey_injected_database_is_still_used(self): + database = unittest.mock.MagicMock() + database.__bool__.return_value = False + database.cleanup_duplicates.return_value = 0 + with ( + unittest.mock.patch( + "py.api.PromptDatabase", side_effect=AssertionError("default DB opened") + ) as factory, + unittest.mock.patch("py.api.get_logger"), + ): + api = PromptManagerAPI(db=database) + self.assertIs(api.db, database) + factory.assert_not_called() + + def test_default_runtime_construction_still_uses_the_configured_factory(self): + database = unittest.mock.MagicMock() + database.cleanup_duplicates.return_value = 2 + with ( + unittest.mock.patch( + "py.api.PromptDatabase", return_value=database + ) as factory, + unittest.mock.patch("py.api.get_logger") as logger, + ): + api = PromptManagerAPI() + self.assertIs(api.db, database) + factory.assert_called_once_with() + logger.return_value.info.assert_any_call( + "Startup cleanup: removed 2 duplicate prompts" + ) + + def test_injected_cleanup_failure_is_logged_without_falling_back_to_default(self): + database = unittest.mock.MagicMock() + database.cleanup_duplicates.side_effect = RuntimeError("test cleanup") + with ( + unittest.mock.patch( + "py.api.PromptDatabase", side_effect=AssertionError("default DB opened") + ) as factory, + unittest.mock.patch("py.api.get_logger") as logger, + ): + api = PromptManagerAPI(db=database) + self.assertIs(api.db, database) + factory.assert_not_called() + logger.return_value.error.assert_called_once_with( + "Startup cleanup failed: test cleanup" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_api_images.py b/tests/test_api_images.py index 2a820ad..c43cedd 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -64,8 +64,7 @@ class ImageAPITestCase(AioHTTPTestCase): app = web.Application() routes = web.RouteTableDef() - self.api = PromptManagerAPI() - self.api.db = PromptDatabase(self._temp_db.name) + self.api = PromptManagerAPI(db=PromptDatabase(self._temp_db.name)) self.api._get_all_output_dirs = lambda: [self.output_dir] self.api._find_comfyui_output_dir = lambda: str(self.output_dir) self.api.add_routes(routes) diff --git a/tests/test_api_logging.py b/tests/test_api_logging.py index cf15f96..97c18c4 100644 --- a/tests/test_api_logging.py +++ b/tests/test_api_logging.py @@ -104,8 +104,9 @@ class LoggingAPITestCase(AioHTTPTestCase): app = web.Application() routes = web.RouteTableDef() - self.api = PromptManagerAPI() - self.api.db = PromptDatabase(os.path.join(self.tmpdir, "prompts.db")) + self.api = PromptManagerAPI( + db=PromptDatabase(os.path.join(self.tmpdir, "prompts.db")) + ) self.api._get_logger_manager = lambda: self.manager self.api.add_routes(routes) app.router.add_routes(routes) diff --git a/tests/test_api_lora.py b/tests/test_api_lora.py index 1c340b1..36b9955 100644 --- a/tests/test_api_lora.py +++ b/tests/test_api_lora.py @@ -76,8 +76,7 @@ class LoraAPITestCase(AioHTTPTestCase): app = web.Application() routes = web.RouteTableDef() - self.api = PromptManagerAPI() - self.api.db = PromptDatabase(self._temp_db.name) + self.api = PromptManagerAPI(db=PromptDatabase(self._temp_db.name)) self.api.add_routes(routes) app.router.add_routes(routes) return app diff --git a/tests/test_api_prompts.py b/tests/test_api_prompts.py index e2a6c83..5773bf7 100644 --- a/tests/test_api_prompts.py +++ b/tests/test_api_prompts.py @@ -38,8 +38,7 @@ class PromptAPITestCase(AioHTTPTestCase): app = web.Application() routes = web.RouteTableDef() - self.api = PromptManagerAPI() - self.api.db = PromptDatabase(self._temp_db.name) + self.api = PromptManagerAPI(db=PromptDatabase(self._temp_db.name)) self.api.add_routes(routes) app.router.add_routes(routes) return app diff --git a/tests/test_gallery_roots.py b/tests/test_gallery_roots.py index 9d1dfef..5ae871b 100644 --- a/tests/test_gallery_roots.py +++ b/tests/test_gallery_roots.py @@ -299,7 +299,9 @@ class TestGetAllOutputDirsFiltersRoots(GalleryRootTestCase): super().setUp() from py.api import PromptManagerAPI - self.api = PromptManagerAPI() + database = MagicMock() + database.cleanup_duplicates.return_value = 0 + self.api = PromptManagerAPI(db=database) self.api._cached_output_dir = None self._orig_dirs = list(GalleryConfig.MONITORING_DIRECTORIES) self.addCleanup(self._restore_dirs) diff --git a/tests/test_output_dir.py b/tests/test_output_dir.py index dd518b0..aebb3a0 100644 --- a/tests/test_output_dir.py +++ b/tests/test_output_dir.py @@ -9,7 +9,7 @@ import sys import tempfile import unittest from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) @@ -26,7 +26,16 @@ class TestFindComfyuiOutputDir(unittest.TestCase): """Test _find_comfyui_output_dir() respects user-configured directories.""" def setUp(self): - self.api = PromptManagerAPI() + roots = patch.object( + GalleryConfig, + "allowed_gallery_parents", + return_value=[tempfile.gettempdir()], + ) + roots.start() + self.addCleanup(roots.stop) + database = MagicMock() + database.cleanup_duplicates.return_value = 0 + self.api = PromptManagerAPI(db=database) self.api._cached_output_dir = None self._orig_dirs = list(GalleryConfig.MONITORING_DIRECTORIES)