diff --git a/__init__.py b/__init__.py index aa7acc6..e33aa42 100644 --- a/__init__.py +++ b/__init__.py @@ -151,6 +151,8 @@ __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] print() print(f"\033[94m[ComfyUI-PromptManager] Version:\033[0m {get_version()}") for node_key, display_name in NODE_DISPLAY_NAME_MAPPINGS.items(): - print(f"🫶 \033[94mLoaded:\033[0m {display_name}") + # ASCII only: a Windows console using a legacy code page cannot print emoji, + # and a UnicodeEncodeError here would abort loading the whole node. + print(f" \033[94mLoaded:\033[0m {display_name}") print(f"\033[94mTotal: {len(NODE_CLASS_MAPPINGS)} tools loaded\033[0m") print() diff --git a/database/models.py b/database/models.py index 3c1d4c7..351835d 100644 --- a/database/models.py +++ b/database/models.py @@ -8,6 +8,8 @@ import shutil import sqlite3 import os import threading +import weakref +from contextlib import closing from pathlib import Path from typing import Set, Tuple @@ -59,6 +61,10 @@ class PromptModel: # same file; only the first one pays for the schema check. _initialized_paths: Set[str] = set() _init_lock = threading.Lock() + # Every live model, so a shutdown or a test teardown can close them all + # before the database file is deleted (Windows refuses otherwise). + _instances: "weakref.WeakSet[PromptModel]" = weakref.WeakSet() + _instances_lock = threading.Lock() def __init__(self, db_path: str): """ @@ -78,6 +84,8 @@ class PromptModel: # Bumped by _close_all_connections so threads holding a closed # connection reopen on their next get_connection(). self._generation = 0 + with PromptModel._instances_lock: + PromptModel._instances.add(self) self._ensure_database_exists() def _ensure_database_exists(self) -> None: @@ -292,6 +300,14 @@ class PromptModel: self._connections.add(conn) return conn + @classmethod + def close_all_instances(cls) -> None: + """Close every connection of every live model in this process.""" + with cls._instances_lock: + models = list(cls._instances) + for model in models: + model.close_all() + def close_all(self) -> None: """Close every thread's connection; they reopen lazily on next use. @@ -627,7 +643,7 @@ class PromptModel: improving query performance and reducing file size. """ try: - with sqlite3.connect(self.db_path) as conn: + with closing(sqlite3.connect(self.db_path)) as conn: conn.execute("VACUUM") conn.commit() except Exception as e: diff --git a/database/operations.py b/database/operations.py index 879c56c..f942a18 100644 --- a/database/operations.py +++ b/database/operations.py @@ -89,7 +89,7 @@ def _resolve_db_path(db_path: Optional[str] = None) -> str: db_path = "prompts.db" if not os.path.isabs(db_path): - db_path = os.path.join(extension_root, db_path) + db_path = os.path.normpath(os.path.join(extension_root, db_path)) # Ensure parent directory exists for custom paths parent = os.path.dirname(db_path) @@ -165,6 +165,11 @@ class PromptDatabase: """Close the calling thread's database connection (see PromptModel.close).""" self.model.close() + @staticmethod + def close_all_instances() -> None: + """Close the connections of every PromptDatabase/PromptModel in the process.""" + PromptModel.close_all_instances() + def close_all(self) -> None: """Close every thread's connection (see PromptModel.close_all).""" self.model.close_all() @@ -667,7 +672,8 @@ class PromptDatabase: rel = os.path.relpath(parent, root) if not rel.startswith(".."): if rel != ".": - folders.add(rel) + # Clients split and filter on "/" whatever the OS + folders.add(rel.replace(os.sep, "/")) made_relative = True break except ValueError: diff --git a/py/api/admin.py b/py/api/admin.py index 5727c30..6a92dc4 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -2,6 +2,7 @@ import asyncio import contextlib +from contextlib import closing import datetime import hashlib import json @@ -190,7 +191,7 @@ def _diagnose_database(db_path): "message": f"Database file not found: {os.path.basename(db_path)}", } try: - with sqlite3.connect(db_path) as conn: + with closing(sqlite3.connect(db_path)) as conn: prompt_count = conn.execute("SELECT COUNT(*) FROM prompts").fetchone()[0] has_images_table = ( conn.execute( @@ -1219,7 +1220,10 @@ class AdminRoutesMixin: removed = await self._run_in_executor(_prune_safety_backups, model.db_path) if removed: self.logger.info(f"Removed {len(removed)} old safety backups") + previous = self.db self.db = PromptDatabase(model.db_path) + # Its connections would otherwise linger (a file lock on Windows) + previous.close_all() prompt_count = self.db.model.get_database_info().get("total_prompts", 0) return web.json_response( diff --git a/requirements-dev.txt b/requirements-dev.txt index 9acf917..7338347 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -5,6 +5,7 @@ aiohttp>=3.8.0 # Tests and coverage coverage>=7.0 +numpy>=1.24 # WD14 engine tests # Code quality black>=24.0.0 diff --git a/tests/open_handles.py b/tests/open_handles.py new file mode 100644 index 0000000..fef1d86 --- /dev/null +++ b/tests/open_handles.py @@ -0,0 +1,43 @@ +"""Teardown guard: fail when a file about to be deleted is still open. + +Windows refuses to delete an open file, so a sqlite connection that was not +closed explicitly turns into a PermissionError in CI there. Linux would let +the deletion through silently; this check makes the Linux run catch the same +mistake. It inspects /proc/self/fd and is a no-op where that does not exist. +""" + +import os + + +def open_handles_under(path): + """Basenames of this process's open files at *path* or below it.""" + fd_dir = "/proc/self/fd" + if not os.path.isdir(fd_dir): + return [] + real = os.path.realpath(str(path)) + hits = [] + for fd in os.listdir(fd_dir): + try: + target = os.readlink(os.path.join(fd_dir, fd)) + except OSError: + continue + if not target.startswith(os.sep): + continue # sockets, pipes, anon inodes + target = target.replace(" (deleted)", "") + if ( + target == real + or target.startswith(real + os.sep) + or target.startswith(real + "-") + ): + hits.append(os.path.basename(target)) + return sorted(hits) + + +def assert_closed(testcase, *paths): + """Fail *testcase* if any of *paths* (files or directories) is still open.""" + hits = [hit for path in paths for hit in open_handles_under(path)] + if hits: + testcase.fail( + "Files still open at teardown (Windows could not delete them): " + + ", ".join(hits) + ) diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index 340ea53..fdb2d1b 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -20,6 +20,12 @@ _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 database.models import PromptModel # noqa: E402 + +try: + from tests.open_handles import assert_closed # noqa: E402 +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed # noqa: E402 from aiohttp import web # noqa: E402 from aiohttp.test_utils import AioHTTPTestCase # noqa: E402 @@ -36,7 +42,9 @@ class AdminAPITestCase(AioHTTPTestCase): """Test app with PromptManager routes, temp DB and config, fake ComfyUI tree.""" async def get_application(self): - self.tmpdir = tempfile.mkdtemp() + # realpath: macOS hands out /var/... for a /private/var/... directory and + # the gallery roots are stored resolved + self.tmpdir = os.path.realpath(tempfile.mkdtemp()) self.comfy_dir = Path(self.tmpdir) / "ComfyUI" self.output_dir = self.comfy_dir / "output" self.output_dir.mkdir(parents=True) @@ -81,7 +89,8 @@ class AdminAPITestCase(AioHTTPTestCase): sys.modules.pop("folder_paths", None) else: sys.modules["folder_paths"] = self._orig_folder_paths - self.api.db.close_all() + PromptModel.close_all_instances() # restore may have swapped api.db + assert_closed(self, self.tmpdir) shutil.rmtree(self.tmpdir, ignore_errors=True) # ── helpers ──────────────────────────────────────────────────────── diff --git a/tests/test_api_backup_restore.py b/tests/test_api_backup_restore.py index 326537a..fb473d9 100644 --- a/tests/test_api_backup_restore.py +++ b/tests/test_api_backup_restore.py @@ -14,6 +14,12 @@ import unittest from unittest.mock import patch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from database.models import PromptModel # noqa: E402 + +try: + from tests.open_handles import assert_closed # noqa: E402 +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed # noqa: E402 import aiohttp from aiohttp import web @@ -47,7 +53,8 @@ class BackupRestoreTestCase(AioHTTPTestCase): async def tearDownAsync(self): await super().tearDownAsync() # closes the aiohttp test client - self.api.db.close_all() + PromptModel.close_all_instances() # restore may have swapped api.db + assert_closed(self, self.tmpdir) for name in os.listdir(self.tmpdir): os.unlink(os.path.join(self.tmpdir, name)) os.rmdir(self.tmpdir) diff --git a/tests/test_api_core.py b/tests/test_api_core.py index 1008cce..16d4a55 100644 --- a/tests/test_api_core.py +++ b/tests/test_api_core.py @@ -16,6 +16,12 @@ sys.modules.setdefault("server", _mock_server) sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +try: + from tests.open_handles import assert_closed # noqa: E402 +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed # noqa: E402 +from database.models import PromptModel # noqa: E402 + from aiohttp import web # noqa: E402 from aiohttp.test_utils import AioHTTPTestCase # noqa: E402 @@ -33,6 +39,8 @@ class FolderPathsFixture(unittest.TestCase): def setUp(self): self.tmpdir = tempfile.mkdtemp() self.addCleanup(shutil.rmtree, self.tmpdir, True) + self.addCleanup(assert_closed, self, self.tmpdir) # just before the rmtree + self.addCleanup(PromptModel.close_all_instances) self.comfy_dir = Path(self.tmpdir) / "ComfyUI" self.output_dir = self.comfy_dir / "output" self.output_dir.mkdir(parents=True) diff --git a/tests/test_api_images.py b/tests/test_api_images.py index b52437d..7c77e31 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -213,6 +213,8 @@ class TestServeOutputImage(ImageAPITestCase): self.assertEqual(hit.status, 200) self.assertEqual(miss.status, 404) self.assertEqual(bad_index.status, 200) + for resp in (hit, miss, bad_index): + await resp.read() # release the served file before the directory goes class TestServeImageById(ImageAPITestCase): @@ -1269,6 +1271,7 @@ class TestLoraManagerPathAllowList(ImageRouteCoverageCase): "GET", f"/prompt_manager/images/{image_id}/file" ) self.assertEqual(resp.status, 200) + await resp.read() # release the served file before the directory goes resolve.assert_called_once_with(configured) diff --git a/tests/test_api_logging.py b/tests/test_api_logging.py index 1614909..cf15f96 100644 --- a/tests/test_api_logging.py +++ b/tests/test_api_logging.py @@ -348,7 +348,7 @@ class TestLogFilesAndDownload(LoggingAPITestCase): self.assertEqual(data["files"], []) async def test_download_existing_file(self): - (self.log_dir / "prompt_manager.log").write_text("line1\nline2\n") + (self.log_dir / "prompt_manager.log").write_bytes(b"line1\nline2\n") resp = await self.client.request( "GET", "/prompt_manager/logs/download/prompt_manager.log" diff --git a/tests/test_basic.py b/tests/test_basic.py index ca70b02..7545848 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -12,6 +12,11 @@ import sys sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +try: + from tests.open_handles import assert_closed # noqa: E402 +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed # noqa: E402 + from database.operations import PromptDatabase from utils.hashing import generate_prompt_hash from utils.validators import validate_prompt_text, validate_rating, validate_tags @@ -30,8 +35,12 @@ class TestBasicFunctionality(unittest.TestCase): def tearDown(self): """Clean up test fixtures.""" # Remove temporary database - if os.path.exists(self.temp_db.name): - os.unlink(self.temp_db.name) + if getattr(self, "db", None) is not None: + self.db.close_all() + assert_closed(self, self.temp_db.name) + for suffix in ("", "-wal", "-shm"): + if os.path.exists(self.temp_db.name + suffix): + os.unlink(self.temp_db.name + suffix) def test_prompt_hash_generation(self): """Test prompt hash generation.""" @@ -214,8 +223,12 @@ class TestNodeIntegration(unittest.TestCase): def tearDown(self): """Clean up test fixtures.""" # Remove temporary database - if os.path.exists(self.temp_db.name): - os.unlink(self.temp_db.name) + if getattr(self, "db", None) is not None: + self.db.close_all() + assert_closed(self, self.temp_db.name) + for suffix in ("", "-wal", "-shm"): + if os.path.exists(self.temp_db.name + suffix): + os.unlink(self.temp_db.name + suffix) @patch("prompt_manager_base.PromptDatabase") def test_node_encode_function(self, mock_db_class): diff --git a/tests/test_database.py b/tests/test_database.py index cd64f6a..b33d233 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -10,6 +10,7 @@ import csv import json import os import sqlite3 +from contextlib import closing import sys import tempfile import types @@ -17,6 +18,12 @@ import unittest from unittest.mock import MagicMock, patch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from database.models import PromptModel # noqa: E402 + +try: + from tests.open_handles import assert_closed # noqa: E402 +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed # noqa: E402 from database.operations import PromptDatabase, _resolve_db_path from utils.hashing import generate_prompt_hash @@ -31,7 +38,8 @@ class DatabaseTestCase(unittest.TestCase): self.db = PromptDatabase(self.temp_db.name) def tearDown(self): - self.db.close_all() + PromptModel.close_all_instances() + assert_closed(self, self.temp_db.name) for suffix in ("", "-wal", "-shm"): path = self.temp_db.name + suffix if os.path.exists(path): @@ -756,7 +764,7 @@ class TestExport(DatabaseTestCase): class TestDuplicateMerging(DatabaseTestCase): def _save_dup(self, text, created_at, **kwargs): pid = self.db.save_prompt(text=text, prompt_hash=None, **kwargs) - with sqlite3.connect(self.temp_db.name) as conn: + with closing(sqlite3.connect(self.temp_db.name)) as conn, conn: conn.execute( "UPDATE prompts SET created_at = ? WHERE id = ?", (created_at, pid) ) @@ -852,7 +860,7 @@ class TestImageQueries(DatabaseTestCase): b = self._save("beta prompt") self.db.link_image_to_prompt(a, "/out/a1.png") self.db.link_image_to_prompt(b, "/out/b1.png") - with sqlite3.connect(self.temp_db.name) as conn: + with closing(sqlite3.connect(self.temp_db.name)) as conn, conn: conn.execute( "UPDATE generated_images SET generation_time = '2026-01-01 00:00:00'" " WHERE filename = 'a1.png'" @@ -938,7 +946,7 @@ class TestImageQueries(DatabaseTestCase): def test_image_row_with_broken_json_yields_empty_dicts(self): pid = self._save("broken") image_id = self.db.link_image_to_prompt(pid, "/out/broken.png") - with sqlite3.connect(self.temp_db.name) as conn: + with closing(sqlite3.connect(self.temp_db.name)) as conn, conn: conn.execute( "UPDATE generated_images SET workflow_data = '{not json', " "prompt_metadata = '{\"k\": 1}', parameters = NULL WHERE id = ?", @@ -1027,7 +1035,7 @@ class TestMaintenanceOperations(DatabaseTestCase): pid = self._save("to orphan", tags=["t"]) self.db.link_image_to_prompt(pid, "/out/o.png") self.db.close() - with sqlite3.connect(self.temp_db.name) as conn: + with closing(sqlite3.connect(self.temp_db.name)) as conn, conn: conn.execute("PRAGMA foreign_keys = OFF") conn.execute("DELETE FROM prompts WHERE id = ?", (pid,)) @@ -1139,7 +1147,7 @@ class TestRowConversion(DatabaseTestCase): def test_search_date_filters(self): pid = self._save("dated") - with sqlite3.connect(self.temp_db.name) as conn: + with closing(sqlite3.connect(self.temp_db.name)) as conn, conn: conn.execute( "UPDATE prompts SET created_at = '2026-05-01T00:00:00' WHERE id = ?", (pid,), diff --git a/tests/test_database_migrations.py b/tests/test_database_migrations.py index 467aa69..73b8011 100644 --- a/tests/test_database_migrations.py +++ b/tests/test_database_migrations.py @@ -16,6 +16,7 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from database.models import PromptModel from database.operations import PromptDatabase +from database.models import normalize_image_path # noqa: E402 LEGACY_PROMPTS_TABLE = ( "CREATE TABLE prompts (id INTEGER PRIMARY KEY AUTOINCREMENT," @@ -260,7 +261,7 @@ class TestIdempotentMigrations(MigrationTestCase): self.assertEqual( sorted(i["file_path"] for i in images), sorted( - os.path.normcase(os.path.normpath(p)) + normalize_image_path(p) for p in ( "/out/2026-01-01/a.png", "/out/2026-01-02/a.png", diff --git a/tests/test_diagnostics_config.py b/tests/test_diagnostics_config.py index 9ecf077..fdbd9c3 100644 --- a/tests/test_diagnostics_config.py +++ b/tests/test_diagnostics_config.py @@ -10,6 +10,7 @@ Verifies that: import asyncio import json import os +from contextlib import closing import sys import tempfile import types @@ -119,7 +120,7 @@ class TestAdminEndpointsUseConfigPath(unittest.TestCase): try: import sqlite3 - with sqlite3.connect(db_file) as conn: + with closing(sqlite3.connect(db_file)) as conn, conn: conn.execute( "CREATE TABLE prompts " "(id INTEGER PRIMARY KEY, text TEXT, created_at TEXT)" diff --git a/tests/test_extension_web_dir.py b/tests/test_extension_web_dir.py index 57f371f..b360f04 100644 --- a/tests/test_extension_web_dir.py +++ b/tests/test_extension_web_dir.py @@ -14,7 +14,7 @@ ROOT = os.path.dirname(TESTS_DIR) sys.path.insert(0, ROOT) sys.path.insert(0, TESTS_DIR) -from test_init_package import load_package # noqa: E402 +from test_init_package import close_every_database, load_package # noqa: E402 class TestExtensionWebDir(unittest.TestCase): @@ -35,6 +35,7 @@ class TestExtensionWebDir(unittest.TestCase): if monitor is not None: monitor.stop_monitoring() sys.modules.pop(cls.package_name, None) + close_every_database() # the package's database must not hold the temp dir open cls.tmp.cleanup() def test_web_directory_is_the_canvas_extension_folder(self): diff --git a/tests/test_init_package.py b/tests/test_init_package.py index b25ee28..b10870a 100644 --- a/tests/test_init_package.py +++ b/tests/test_init_package.py @@ -17,6 +17,11 @@ from aiohttp import web ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, ROOT) +try: + from tests.open_handles import open_handles_under # noqa: E402 +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import open_handles_under # noqa: E402 + class FakePromptQueue: def __init__(self): @@ -82,6 +87,17 @@ def load_package(name, db_path, output_dir, server=None): return package, server.PromptServer.instance +def close_every_database(): + """Close connections of every PromptModel registry in the process. + + A package loaded under its own name imports its own copy of + database.models, whose PromptModel class keeps its own registry. + """ + for name, module in list(sys.modules.items()): + if name.endswith("database.models") and hasattr(module, "PromptModel"): + module.PromptModel.close_all_instances() + + class PackageTestCase(unittest.TestCase): package_name = "prompt_manager_pkg" @@ -109,7 +125,10 @@ class PackageTestCase(unittest.TestCase): if monitor is not None: monitor.stop_monitoring() sys.modules.pop(cls.package_name, None) + close_every_database() + hits = open_handles_under(cls.tmp.name) cls.tmp.cleanup() + assert not hits, f"files still open before temp dir cleanup: {hits}" class TestPackageWiring(PackageTestCase): @@ -220,3 +239,30 @@ class TestPackageWhenHookRegistrationFails(PackageTestCase): if __name__ == "__main__": unittest.main() + + +class TestConsoleBanner(unittest.TestCase): + """Loading the node must not depend on the console being able to print emoji.""" + + def test_package_loads_with_a_cp1252_console(self): + import contextlib + import io + + tmp = tempfile.TemporaryDirectory() + self.addCleanup(tmp.cleanup) + self.addCleanup(close_every_database) + output_dir = os.path.join(tmp.name, "output") + os.makedirs(output_dir) + name = "prompt_manager_pkg_cp1252" + stream = io.TextIOWrapper(io.BytesIO(), encoding="cp1252", errors="strict") + + with contextlib.redirect_stdout(stream): + package, _ = load_package(name, os.path.join(tmp.name, "p.db"), output_dir) + self.addCleanup(sys.modules.pop, name, None) + monitor = getattr(package, "_global_image_monitor", None) + if monitor is not None: + self.addCleanup(monitor.stop_monitoring) + + stream.flush() + banner = stream.buffer.getvalue().decode("cp1252") + self.assertIn("Loaded:", banner) diff --git a/tests/test_lora_database.py b/tests/test_lora_database.py index 5be8c9a..648b821 100644 --- a/tests/test_lora_database.py +++ b/tests/test_lora_database.py @@ -195,6 +195,25 @@ class TestGetPromptSubfolders(LoraDBTestCase): self.assertIsInstance(folders, list) self.assertGreater(len(folders), 0) + def test_subfolders_use_forward_slashes_on_every_platform(self): + """Clients split folders on "/" whatever separator the OS produced.""" + from unittest import mock + + import database.operations as operations + + pid = self._save("prompt") + self._link_image(pid, "/output/2026/08-Aug/06/img.png") + + with ( + mock.patch.object( + operations.os.path, "relpath", return_value="2026\\08-Aug\\06" + ), + mock.patch.object(operations.os, "sep", "\\"), + ): + folders = self.db.get_prompt_subfolders(root_dirs=["/output"]) + + self.assertEqual(folders, ["2026/08-Aug/06"]) + def test_include_ancestors_adds_intermediate_paths(self): pid = self._save("prompt") self._link_image(pid, "/output/2026/08-Aug/2026-08-06/img.png") diff --git a/tests/test_lora_utils.py b/tests/test_lora_utils.py index 61ab15b..b273643 100644 --- a/tests/test_lora_utils.py +++ b/tests/test_lora_utils.py @@ -229,7 +229,7 @@ class TestGetLoraImageCacheDir(unittest.TestCase): def test_returns_path(self): cache_dir = get_lora_image_cache_dir() self.assertIsInstance(cache_dir, Path) - self.assertTrue(str(cache_dir).endswith("data/lora_images")) + self.assertEqual(cache_dir.parts[-2:], ("data", "lora_images")) def test_directory_exists(self): cache_dir = get_lora_image_cache_dir() diff --git a/tests/test_prompt_usage.py b/tests/test_prompt_usage.py index fd5fa19..c460664 100644 --- a/tests/test_prompt_usage.py +++ b/tests/test_prompt_usage.py @@ -2,12 +2,18 @@ import os import sqlite3 +from contextlib import closing import sys import tempfile import unittest sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +try: + from tests.open_handles import assert_closed # noqa: E402 +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed # noqa: E402 + from database.models import PromptModel from database.operations import PromptDatabase from prompt_manager_base import PromptManagerBase @@ -23,6 +29,8 @@ class UsageTestCase(unittest.TestCase): def tearDown(self): if getattr(self, "db", None) is not None: self.db.close_all() + PromptModel.close_all_instances() + assert_closed(self, self.path) for suffix in ("", "-wal", "-shm"): if os.path.exists(self.path + suffix): os.unlink(self.path + suffix) @@ -33,7 +41,7 @@ class UsageTestCase(unittest.TestCase): ) def _set(self, prompt_id, **columns): - with sqlite3.connect(self.path) as conn: + with closing(sqlite3.connect(self.path)) as conn, conn: for column, value in columns.items(): conn.execute( f"UPDATE prompts SET {column} = ? WHERE id = ?", (value, prompt_id) @@ -169,7 +177,7 @@ class TestUsageMigration(unittest.TestCase): self.with_images = db.save_prompt(text="has images", prompt_hash="h1") self.no_images = db.save_prompt(text="no images", prompt_hash="h2") db.close_all() - with sqlite3.connect(self.path) as conn: + with closing(sqlite3.connect(self.path)) as conn, conn: conn.execute( "UPDATE prompts SET created_at = '2026-01-01T10:00:00.000000+00:00'" ) @@ -199,6 +207,8 @@ class TestUsageMigration(unittest.TestCase): def tearDown(self): for db in self.dbs: db.close_all() + PromptModel.close_all_instances() + assert_closed(self, self.path) for suffix in ("", "-wal", "-shm"): if os.path.exists(self.path + suffix): os.unlink(self.path + suffix) @@ -218,7 +228,7 @@ class TestUsageMigration(unittest.TestCase): # A row whose usage is unknown (e.g. written by 3.2.3 after a downgrade) # stays "never run": a restart must not invent a run for it, otherwise # Recently Used could never tell saved-but-unused prompts from run ones. - with sqlite3.connect(self.path) as conn: + with closing(sqlite3.connect(self.path)) as conn, conn: conn.execute( "INSERT INTO prompts (text, hash, created_at)" " VALUES ('from 3.2.3', 'h3', '2026-04-01T00:00:00.000000+00:00')" diff --git a/tests/test_utils_diagnostics.py b/tests/test_utils_diagnostics.py index 33353d8..c001233 100644 --- a/tests/test_utils_diagnostics.py +++ b/tests/test_utils_diagnostics.py @@ -3,6 +3,7 @@ import os import pathlib import sqlite3 +from contextlib import closing import sys import tempfile import types @@ -75,7 +76,7 @@ class TestDatabaseChecks(DiagnosticsTestCase): self.assertEqual(self.diag.check_images_table()["status"], "error") def test_database_without_images_table(self): - with sqlite3.connect(self.db_path) as conn: + with closing(sqlite3.connect(self.db_path)) as conn, conn: conn.execute("CREATE TABLE prompts (id INTEGER PRIMARY KEY, text TEXT)") self.assertEqual( self.diag.check_database(), diff --git a/utils/diagnostics.py b/utils/diagnostics.py index 14f1581..c0c72d4 100644 --- a/utils/diagnostics.py +++ b/utils/diagnostics.py @@ -32,6 +32,7 @@ The diagnostics provide: import os import sqlite3 +from contextlib import closing from pathlib import Path from typing import Dict, Any @@ -147,7 +148,7 @@ class GalleryDiagnostics: "message": f"Database file not found: {self.db_path}", } - with sqlite3.connect(self.db_path) as conn: + with closing(sqlite3.connect(self.db_path)) as conn: conn.row_factory = sqlite3.Row # Check prompts table @@ -189,7 +190,7 @@ class GalleryDiagnostics: self.logger.info("\n[IMG] Checking Images Table...") try: - with sqlite3.connect(self.db_path) as conn: + with closing(sqlite3.connect(self.db_path)) as conn: conn.row_factory = sqlite3.Row # Check if table exists