Files
ComfyAssets-ComfyUI_PromptM…/tests/test_lora_database.py
T
Vito d88cf9e7c6 fix(quality): sweep the Code Quality standard findings (#558)
* fix(quality): preserve cursor iteration and option validation

Keep prefetched rows on ManagedCursor so partial loops and explicit fetches share one position, and stop existing iterators after close or retirement. Preserve chunked reads and native SQLite result lifetime on reexecution and errors.

Catch OptionError before ValueError so legacy auto-tag requests receive the actual validation error. Remove the overwritten ViewerJS title option.

Regression tests failed before the fixes for lost rows, reads after close/retirement, and misleading option errors. Focused Python checks pass; cursor checks also pass on Python 3.13 and 3.14. Refs #321.

* fix(quality): make the test doubles legal raises and hashes

The fake queue no longer raises an attribute that starts as None. The
oversized kind records hashes instead of raising from __hash__, which
also catches a hash that production code swallows. Hostile.__bool__
raises the TypeError a truth test may raise.

* fix(quality): drop the extra saveCoverage argument

* refactor(quality): remove assignments that are overwritten before use

* refactor(quality): remove the remaining overwritten test assignments

* refactor(quality): drop imports repeated inside functions

* refactor(quality): give duplicate identity keys one shape

* refactor(quality): name the unused job event loop variable

* refactor(quality): release model gates in finally, not except BaseException

* refactor(quality): end the job runners with an explicit return

* refactor(quality): use the precise unittest asserts

* refactor(quality): import unittest.mock one way in the tests

* fix(quality): say why each empty except is safe, log the broad ones

* refactor(quality): stop binding values the tests never read

* refactor(quality): remove the unused preview aliases and tag page query

* refactor(quality): remove re-exports nothing imports

* refactor(quality): parenthesize deliberate string joins in lists

* refactor(quality): remove unused JavaScript variables and imports

* refactor(quality): use the precise assert in the model status test
2026-10-31 05:43:48 +00:00

305 lines
11 KiB
Python

"""
Database tests for LoRA Manager integration and folder filter features.
Tests delete_prompts_by_category, search_prompts folder filter,
get_prompt_subfolders, and LoRA-specific prompt workflows using
an in-memory SQLite database.
"""
import os
import sys
import tempfile
import unittest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database.operations import PromptDatabase
from utils.hashing import generate_prompt_hash
class LoraDBTestCase(unittest.TestCase):
"""Base class with temp database setup/teardown."""
def setUp(self):
self.temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
self.temp_db.close()
self.db = PromptDatabase(self.temp_db.name)
def tearDown(self):
self.db.close()
for suffix in ("", "-wal", "-shm"):
path = self.temp_db.name + suffix
if os.path.exists(path):
os.unlink(path)
def _save(self, text, category=None, tags=None):
"""Save a prompt and return its ID."""
return self.db.save_prompt(
text=text,
category=category,
tags=tags or [],
prompt_hash=generate_prompt_hash(text),
)
def _link_image(self, prompt_id, image_path):
"""Link a fake image to a prompt."""
return self.db.link_image_to_prompt(
prompt_id=str(prompt_id), image_path=image_path
)
# ── delete_prompts_by_category ────────────────────────────────────────
class TestDeleteByCategory(LoraDBTestCase):
"""Test delete_prompts_by_category for LoRA reimport cleanup."""
def test_deletes_matching_category(self):
self._save("lora prompt 1", category="lora-manager")
self._save("lora prompt 2", category="lora-manager")
self._save("keep this", category="general")
deleted = self.db.delete_prompts_by_category("lora-manager")
self.assertEqual(deleted, 2)
results = self.db.search_prompts(category="lora-manager")
self.assertEqual(len(results), 0)
def test_preserves_other_categories(self):
self._save("keep this", category="general")
self._save("and this", category="portraits")
self.db.delete_prompts_by_category("lora-manager")
results = self.db.search_prompts()
self.assertEqual(len(results), 2)
def test_returns_zero_when_none_match(self):
self._save("no match", category="general")
deleted = self.db.delete_prompts_by_category("lora-manager")
self.assertEqual(deleted, 0)
def test_cascades_to_images(self):
pid = self._save("lora with image", category="lora-manager")
self._link_image(pid, "/fake/path/image.jpg")
# Verify image is linked
images = self.db.get_prompt_images(pid)
self.assertEqual(len(images), 1)
self.db.delete_prompts_by_category("lora-manager")
# Prompt gone
results = self.db.search_prompts(category="lora-manager")
self.assertEqual(len(results), 0)
def test_empty_category_string(self):
self._save("test", category="general")
deleted = self.db.delete_prompts_by_category("")
self.assertEqual(deleted, 0)
# ── Folder filter (search_prompts with folder param) ──────────────────
class TestFolderFilter(LoraDBTestCase):
"""Test search_prompts folder parameter for subfolder filtering."""
def _setup_prompts_with_images(self):
"""Create prompts linked to images in different directories."""
pid1 = self._save("landscape prompt", category="nature")
self._link_image(pid1, "/output/landscapes/sunset.png")
pid2 = self._save("portrait prompt", category="portraits")
self._link_image(pid2, "/output/portraits/face.png")
pid3 = self._save("another landscape", category="nature")
self._link_image(pid3, "/output/landscapes/mountain.png")
return pid1, pid2, pid3
def test_filter_by_folder(self):
self._setup_prompts_with_images()
results = self.db.search_prompts(folder="landscapes")
self.assertEqual(len(results), 2)
texts = {r["text"] for r in results}
self.assertEqual(texts, {"landscape prompt", "another landscape"})
def test_filter_different_folder(self):
self._setup_prompts_with_images()
results = self.db.search_prompts(folder="portraits")
self.assertEqual(len(results), 1)
self.assertEqual(results[0]["text"], "portrait prompt")
def test_no_match_returns_empty(self):
self._setup_prompts_with_images()
results = self.db.search_prompts(folder="nonexistent")
self.assertEqual(len(results), 0)
def test_no_folder_returns_all(self):
self._setup_prompts_with_images()
results = self.db.search_prompts()
self.assertGreaterEqual(len(results), 3)
def test_folder_with_category_filter(self):
self._setup_prompts_with_images()
results = self.db.search_prompts(folder="landscapes", category="nature")
self.assertEqual(len(results), 2)
# ── get_prompt_subfolders ─────────────────────────────────────────────
class TestGetPromptSubfolders(LoraDBTestCase):
"""Test get_prompt_subfolders — extracts unique folder names from images."""
def test_extracts_subfolders(self):
pid1 = self._save("prompt 1")
self._link_image(pid1, "/output/folder_a/img1.png")
pid2 = self._save("prompt 2")
self._link_image(pid2, "/output/folder_b/img2.png")
folders = self.db.get_prompt_subfolders()
self.assertIsInstance(folders, list)
self.assertGreaterEqual(len(folders), 2)
def test_deduplicates(self):
pid1 = self._save("prompt 1")
self._link_image(pid1, "/output/same_folder/img1.png")
pid2 = self._save("prompt 2")
self._link_image(pid2, "/output/same_folder/img2.png")
folders = self.db.get_prompt_subfolders()
# Count occurrences of the folder — should appear once
matches = [f for f in folders if "same_folder" in f]
self.assertEqual(len(matches), 1)
def test_empty_database(self):
folders = self.db.get_prompt_subfolders()
self.assertEqual(folders, [])
def test_returns_sorted(self):
for i, name in enumerate(["charlie", "alpha", "bravo"]):
pid = self._save(f"prompt {i}")
self._link_image(pid, f"/output/{name}/img.png")
folders = self.db.get_prompt_subfolders()
self.assertEqual(folders, sorted(folders))
def test_with_root_dirs(self):
pid = self._save("prompt")
self._link_image(pid, "/output/sub/deep/img.png")
folders = self.db.get_prompt_subfolders(root_dirs=["/output"])
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."""
import unittest.mock as 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")
folders = self.db.get_prompt_subfolders(
root_dirs=["/output"], include_ancestors=True
)
self.assertIn("2026", folders)
self.assertIn("2026/08-Aug", folders)
self.assertIn("2026/08-Aug/2026-08-06", folders)
def test_include_ancestors_false_no_intermediates(self):
pid = self._save("prompt")
self._link_image(pid, "/output/2026/08-Aug/2026-08-06/img.png")
folders = self.db.get_prompt_subfolders(
root_dirs=["/output"], include_ancestors=False
)
self.assertIn("2026/08-Aug/2026-08-06", folders)
self.assertNotIn("2026", folders)
self.assertNotIn("2026/08-Aug", folders)
def test_include_ancestors_sorted(self):
pid = self._save("prompt")
self._link_image(pid, "/output/a/b/c/img.png")
folders = self.db.get_prompt_subfolders(
root_dirs=["/output"], include_ancestors=True
)
self.assertEqual(folders, sorted(folders))
# ── LoRA prompt workflow ──────────────────────────────────────────────
class TestLoraPromptWorkflow(LoraDBTestCase):
"""Test the full LoRA import workflow at the database layer."""
def test_save_lora_prompt_with_tags(self):
"""Simulate what lora_scan does: save prompt with lora-manager tags."""
pid = self._save(
text="1girl, detailed face, anime style",
category="lora-manager",
tags=["lora-manager", "lora:my_lora", "trigger1"],
)
prompt = self.db.get_prompt_by_id(pid)
self.assertEqual(prompt["category"], "lora-manager")
self.assertIn("lora-manager", prompt["tags"])
def test_reimport_clears_and_recreates(self):
"""Simulate reimport: delete old, create new."""
# First import
pid1 = self._save("old lora prompt", category="lora-manager")
self._link_image(pid1, "/cache/old.jpg")
# Reimport
self.db.delete_prompts_by_category("lora-manager")
# Second import
pid2 = self._save("new lora prompt", category="lora-manager")
self._link_image(pid2, "/cache/new.jpg")
results = self.db.search_prompts(category="lora-manager")
self.assertEqual(len(results), 1)
self.assertEqual(results[0]["text"], "new lora prompt")
def test_hash_dedup_prevents_duplicates(self):
"""Verify hash-based dedup works for LoRA prompts."""
text = "duplicate lora prompt"
h = generate_prompt_hash(text)
self._save(text, category="lora-manager")
existing = self.db.get_prompt_by_hash(h)
self.assertIsNotNone(existing)
def test_link_multiple_images_to_lora_prompt(self):
"""LoRA prompts can have multiple preview images."""
pid = self._save("multi-image lora", category="lora-manager")
self._link_image(pid, "/cache/lora/img1.jpg")
self._link_image(pid, "/cache/lora/img2.jpg")
self._link_image(pid, "/cache/lora/img3.jpg")
images = self.db.get_prompt_images(pid)
self.assertEqual(len(images), 3)
if __name__ == "__main__":
unittest.main()