diff --git a/tests/private_temp.py b/tests/private_temp.py new file mode 100644 index 00000000..dee79ce9 --- /dev/null +++ b/tests/private_temp.py @@ -0,0 +1,24 @@ +"""A temp SQLite database in a folder no other test process shares (#346). + +``PromptManagerAPI`` sweeps its database's folder at startup for the stages, +uploads and safety-backup copies an interrupted restore or backup left. A +fixture database directly in the system temp folder made every other test +process that started an API (another worktree's suite, a pre-push hook) delete +this process's in-flight restore files, and left safety backups behind there. +""" + +import tempfile + + +def private_temp_db(case): + """Return a closed ``NamedTemporaryFile`` handle for an empty ``.db`` file. + + The file lives in a new folder that ``case`` removes at cleanup, after its + own teardown has closed and unlinked the database (cleanups run after + ``tearDown``), together with anything a restore or backup left next to it. + """ + folder = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) + case.addCleanup(folder.cleanup) + handle = tempfile.NamedTemporaryFile(delete=False, suffix=".db", dir=folder.name) + handle.close() + return handle diff --git a/tests/test_api.py b/tests/test_api.py index 051881a8..efcb433a 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -8,7 +8,6 @@ without requiring a full ComfyUI server. import json import os import sys -import tempfile import unittest sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) @@ -18,6 +17,7 @@ from aiohttp.test_utils import AioHTTPTestCase from database.operations import PromptDatabase from py.api import PromptManagerAPI +from tests.private_temp import private_temp_db from utils.hashing import generate_prompt_hash @@ -25,8 +25,7 @@ class APITestCase(AioHTTPTestCase): """Base class that stands up a test aiohttp app with PromptManager routes.""" async def get_application(self): - self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db") - self._temp_db.close() + self._temp_db = private_temp_db(self) app = web.Application() routes = web.RouteTableDef() diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index e2d2dffc..cfe3f47a 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -35,6 +35,7 @@ from aiohttp.test_utils import AioHTTPTestCase from database.operations import PromptDatabase from py.api import PromptManagerAPI +from tests.private_temp import private_temp_db from py.api.autotag_routes import MODEL_BUSY_ERROR from py.api.jobs import JobBusy from utils.hashing import generate_prompt_hash @@ -143,8 +144,7 @@ class AutotagAPITestCase(AioHTTPTestCase): fake engine.""" async def get_application(self): - self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db") - self._temp_db.close() + self._temp_db = private_temp_db(self) self.tmp_root = Path(tempfile.mkdtemp()) self.output_dir = self.tmp_root / "output" self.output_dir.mkdir() diff --git a/tests/test_api_history.py b/tests/test_api_history.py index 67476d7a..dc952332 100644 --- a/tests/test_api_history.py +++ b/tests/test_api_history.py @@ -1169,3 +1169,47 @@ class HistoryHostCleanupTests(unittest.IsolatedAsyncioTestCase): fence.validate(epoch, generation) with self.assertRaises(sqlite3.Error): api._history_fence() + + +class SharedTempFolderRestoreTests(ImageAPITestCase): + """#346: a restore survives another suite starting in the shared temp folder. + + Every ``PromptManagerAPI`` sweeps its database's folder for stages and + uploads an interrupted restore left. A fixture database directly in the + shared temp folder let any other test process (another worktree's suite, + a pre-push hook) delete this restore's files mid-request. The shared + folder here is a scratch stand-in for the system temp folder, so the test + never sweeps the real one. + """ + + async def get_application(self): + shared = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) + # Cleanups run last-in first-out: this one after the fixture's own. + self.addCleanup(shared.cleanup) + self.shared_tmp = shared.name + with patch.object(tempfile, "tempdir", shared.name): + return await super().get_application() + + async def test_restore_survives_another_suite_sweeping_the_shared_folder(self): + from database import restore_stage + from database.models import PromptModel + + self.api.db.save_prompt("kept by the upload") + source = self.comfy_dir / "upload.db" + self.assertTrue(self.api.db.model.backup_database(str(source))) + stage = PromptModel.stage_restore + swept = [] + + def stage_then_other_suite_starts(model, src_path): + staged = stage(model, src_path) + # What another suite's PromptManagerAPI does at startup. + other = os.path.join(self.shared_tmp, "other-suite.db") + swept.extend(restore_stage.sweep_restore_temps(other)) + return staged + + form = FormData() + form.add_field("database_file", source.read_bytes(), filename="upload.db") + with patch.object(PromptModel, "stage_restore", stage_then_other_suite_starts): + response = await self.client.post("/prompt_manager/restore", data=form) + self.assertEqual(response.status, 200, await response.text()) + self.assertEqual(swept, [], "another suite's sweep reached this restore") diff --git a/tests/test_api_images.py b/tests/test_api_images.py index 08713ddb..cc74c6cf 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -42,6 +42,7 @@ from py import lora_utils # noqa: E402 from py import lora_manager_layout as layout_module # noqa: E402 from py.api.thumbnail_settings import ThumbnailSettings # noqa: E402 from py.api import PromptManagerAPI # noqa: E402 +from tests.private_temp import private_temp_db # noqa: E402 from py.api.prompts import MAX_PAGE_LIMIT # noqa: E402 from py.config import GalleryConfig # noqa: E402 from utils.hashing import generate_prompt_hash # noqa: E402 @@ -60,8 +61,7 @@ class ImageAPITestCase(AioHTTPTestCase): """App with PromptManager routes, a temp DB and a temp output directory.""" async def get_application(self): - self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db") - self._temp_db.close() + self._temp_db = private_temp_db(self) self._output_tmp = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) # A fake ComfyUI tree: /output is the gallery root and is the diff --git a/tests/test_api_lora.py b/tests/test_api_lora.py index 8fda0562..99446360 100644 --- a/tests/test_api_lora.py +++ b/tests/test_api_lora.py @@ -36,6 +36,7 @@ from aiohttp.test_utils import AioHTTPTestCase from database.operations import PromptDatabase from py.api import PromptManagerAPI +from tests.private_temp import private_temp_db from py import config_store, lora_import, lora_utils from py.config import IntegrationConfig, PromptManagerConfig from py.api import settings_routes as settings_module @@ -57,8 +58,7 @@ class LoraAPITestCase(AioHTTPTestCase): """Stands up the PromptManager routes with a temp DB and neutral LoRA config.""" async def get_application(self): - self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db") - self._temp_db.close() + self._temp_db = private_temp_db(self) self.tmp_root = Path(tempfile.mkdtemp()) self.custom_nodes = self.tmp_root / "custom_nodes" self.custom_nodes.mkdir() diff --git a/tests/test_api_prompts.py b/tests/test_api_prompts.py index c4af9c68..c8625c25 100644 --- a/tests/test_api_prompts.py +++ b/tests/test_api_prompts.py @@ -5,7 +5,6 @@ Route tests for the prompt API (py/api/prompts.py). import json import os import sys -import tempfile import unittest from unittest.mock import MagicMock, patch @@ -29,6 +28,7 @@ except ImportError: # discovered with tests/ as the top-level directory from database.operations import PromptDatabase # noqa: E402 from py.api import PromptManagerAPI # noqa: E402 +from tests.private_temp import private_temp_db # noqa: E402 import py.api.prompts as prompts_module # noqa: E402 from py.api.prompts import ( # noqa: E402 BULK_CATEGORY_MAX_PROMPTS, @@ -43,8 +43,7 @@ class PromptAPITestCase(AioHTTPTestCase): """App with PromptManager routes and a temporary SQLite database.""" async def get_application(self): - self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db") - self._temp_db.close() + self._temp_db = private_temp_db(self) app = web.Application() routes = web.RouteTableDef()