nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: first CI run on Windows and macOS; tests now enforce closed database files

Browse Source 
The first push of the branch ran the Tests matrix on all three operating
systems. Every Python job failed; the causes were:

- numpy missing in CI: the WD14 engine tests import it (requirements-dev)
- macOS: temp directories resolve through /private; the admin API tests
  now use the real path like the gallery roots they compare against
- Windows: the startup banner printed an emoji, which raises
  UnicodeEncodeError on a legacy-code-page console and aborted loading
  the whole node; the banner is ASCII now and a test loads the package
  under a cp1252 stdout
- Windows: gallery subfolders came back with backslashes, which neither
  the folder tree nor the search filter understand; they are "/" joined
  on every platform. A relative DEFAULT_DB_PATH kept its forward slashes
  after being joined to the extension root; it is normalised now
- Windows: sqlite connections opened with `with sqlite3.connect()` are
  not closed by the context manager and linger until garbage collection,
  which holds a file lock there. Production sites (vacuum, diagnostics,
  restore inspection) close explicitly; a restore that swaps the API's
  database object closes the previous one; PromptModel keeps a registry
  so a shutdown or a test can close every instance
- Windows: served files were still open when their temp directory was
  removed (tests now read the body), a log download test compared text
  mode newlines, two tests compared paths with a hard-coded separator

tests/open_handles.py makes the Linux run fail the way Windows would:
every database test base checks, right before deleting its files, that
nothing in the process still holds them open.
This commit is contained in:
Vito Sansevero
2026-10-08 18:24:18 -07:00
parent 08d57b408a
commit b1bfce702d
22 changed files with 229 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+3 -1
View File
@@ -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()
database/models.py
+17 -1
View File
@@ -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:
database/operations.py
+8 -2
View File
@@ -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:
py/api/admin.py
+5 -1
View File
@@ -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(
requirements-dev.txt
+1
View File
@@ -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
tests/open_handles.py
+43
View File
@@ -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)
)
tests/test_api_admin.py
+11 -2
View File
@@ -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 ────────────────────────────────────────────────────────
tests/test_api_backup_restore.py
+8 -1
View File
@@ -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)
tests/test_api_core.py
+8
View File
@@ -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)
tests/test_api_images.py
+3
View File
@@ -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)
tests/test_api_logging.py
+1 -1
View File
@@ -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"
tests/test_basic.py
+17 -4
View File
@@ -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):
tests/test_database.py
+14 -6
View File
@@ -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,),
tests/test_database_migrations.py
+2 -1
View File
@@ -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",
tests/test_diagnostics_config.py
+2 -1
View File
@@ -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)"
tests/test_extension_web_dir.py
+2 -1
View File
@@ -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):
tests/test_init_package.py
+46
View File
@@ -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)
tests/test_lora_database.py
+19
View File
@@ -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")
tests/test_lora_utils.py
+1 -1
View File
@@ -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()
tests/test_prompt_usage.py
+13 -3
View File
@@ -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')"
tests/test_utils_diagnostics.py
+2 -1
View File
@@ -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(),
utils/diagnostics.py
+3 -2
View File
@@ -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