test: add unit tests for LoRA integration and folder filter (#52)
- test_lora_utils.py: 32 tests covering trigger word extraction, example prompts, image URLs, metadata parsing, cache dir, TriggerWordCache - test_lora_database.py: 17 tests covering delete_prompts_by_category, folder filter search, get_prompt_subfolders, LoRA import workflow - test_config.py: 7 new IntegrationConfig tests for structure, enable, partial update, reset, and roundtrip
This commit is contained in:
Vito Sansevero
+85
-1
@@ -19,7 +19,7 @@ sys.modules["server"] = _mock_server
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from py.config import GalleryConfig, PromptManagerConfig
|
||||
from py.config import GalleryConfig, IntegrationConfig, PromptManagerConfig
|
||||
|
||||
|
||||
class TestGalleryConfig(unittest.TestCase):
|
||||
@@ -207,5 +207,89 @@ class TestPromptManagerConfig(unittest.TestCase):
|
||||
os.unlink(tmp.name)
|
||||
|
||||
|
||||
class TestIntegrationConfig(unittest.TestCase):
|
||||
"""Test IntegrationConfig for LoRA Manager settings."""
|
||||
|
||||
def setUp(self):
|
||||
"""Save original values to restore after each test."""
|
||||
self._orig = IntegrationConfig.get_config()
|
||||
|
||||
def tearDown(self):
|
||||
"""Restore original config values."""
|
||||
IntegrationConfig.update_config(self._orig)
|
||||
|
||||
def test_reset_to_disabled(self):
|
||||
"""Integrations can be fully disabled via update_config."""
|
||||
IntegrationConfig.update_config(
|
||||
{
|
||||
"lora_manager": {
|
||||
"enabled": False,
|
||||
"path": "",
|
||||
"trigger_words_enabled": False,
|
||||
"civitai_api_key": "",
|
||||
}
|
||||
}
|
||||
)
|
||||
config = IntegrationConfig.get_config()
|
||||
lora = config["lora_manager"]
|
||||
self.assertFalse(lora["enabled"])
|
||||
self.assertEqual(lora["path"], "")
|
||||
self.assertFalse(lora["trigger_words_enabled"])
|
||||
self.assertEqual(lora["civitai_api_key"], "")
|
||||
|
||||
def test_get_config_structure(self):
|
||||
config = IntegrationConfig.get_config()
|
||||
self.assertIn("lora_manager", config)
|
||||
lora = config["lora_manager"]
|
||||
self.assertIn("enabled", lora)
|
||||
self.assertIn("path", lora)
|
||||
self.assertIn("trigger_words_enabled", lora)
|
||||
self.assertIn("civitai_api_key", lora)
|
||||
|
||||
def test_update_config_enables(self):
|
||||
IntegrationConfig.update_config(
|
||||
{
|
||||
"lora_manager": {
|
||||
"enabled": True,
|
||||
"path": "/some/path",
|
||||
"trigger_words_enabled": True,
|
||||
"civitai_api_key": "test-key-123",
|
||||
}
|
||||
}
|
||||
)
|
||||
self.assertTrue(IntegrationConfig.LORA_MANAGER_ENABLED)
|
||||
self.assertEqual(IntegrationConfig.LORA_MANAGER_PATH, "/some/path")
|
||||
self.assertTrue(IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED)
|
||||
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, "test-key-123")
|
||||
|
||||
def test_update_partial(self):
|
||||
"""Updating one field shouldn't affect others."""
|
||||
# Reset to known state first
|
||||
IntegrationConfig.update_config(
|
||||
{"lora_manager": {"enabled": False, "path": "/known"}}
|
||||
)
|
||||
# Now update only enabled
|
||||
IntegrationConfig.update_config({"lora_manager": {"enabled": True}})
|
||||
self.assertTrue(IntegrationConfig.LORA_MANAGER_ENABLED)
|
||||
self.assertEqual(IntegrationConfig.LORA_MANAGER_PATH, "/known")
|
||||
|
||||
def test_update_empty_dict_noop(self):
|
||||
"""Updating with empty dict preserves current state."""
|
||||
before = IntegrationConfig.get_config()
|
||||
IntegrationConfig.update_config({})
|
||||
after = IntegrationConfig.get_config()
|
||||
self.assertEqual(before, after)
|
||||
|
||||
def test_roundtrip(self):
|
||||
"""get_config → update_config → get_config should be stable."""
|
||||
IntegrationConfig.update_config(
|
||||
{"lora_manager": {"enabled": True, "path": "/test"}}
|
||||
)
|
||||
config1 = IntegrationConfig.get_config()
|
||||
IntegrationConfig.update_config(config1)
|
||||
config2 = IntegrationConfig.get_config()
|
||||
self.assertEqual(config1, config2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,253 @@
|
||||
"""
|
||||
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):
|
||||
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)
|
||||
|
||||
|
||||
# ── 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()
|
||||
@@ -0,0 +1,316 @@
|
||||
"""
|
||||
Unit tests for LoRA Manager integration utilities.
|
||||
|
||||
Tests metadata parsing, trigger word extraction, image URL extraction,
|
||||
directory detection, TriggerWordCache, and image download logic.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from py.lora_utils import (
|
||||
TriggerWordCache,
|
||||
get_civitai_image_urls,
|
||||
get_example_prompt_from_metadata,
|
||||
get_lora_image_cache_dir,
|
||||
get_trigger_words_from_metadata,
|
||||
read_lora_metadata,
|
||||
)
|
||||
|
||||
# ── Sample metadata fixtures ───────────────────────────────────────────
|
||||
|
||||
|
||||
def _make_metadata(
|
||||
trained_words=None,
|
||||
images=None,
|
||||
model_name="test_lora",
|
||||
file_name="test.safetensors",
|
||||
):
|
||||
"""Build a realistic LoRA metadata dict for testing."""
|
||||
meta = {"file_name": file_name}
|
||||
civitai = {}
|
||||
if trained_words is not None:
|
||||
civitai["trainedWords"] = trained_words
|
||||
if images is not None:
|
||||
civitai["images"] = images
|
||||
if model_name:
|
||||
civitai["model"] = {"name": model_name}
|
||||
if civitai:
|
||||
meta["civitai"] = civitai
|
||||
return meta
|
||||
|
||||
|
||||
# ── Pure function tests (no mocking) ──────────────────────────────────
|
||||
|
||||
|
||||
class TestGetTriggerWords(unittest.TestCase):
|
||||
"""Test get_trigger_words_from_metadata — pure dict extraction."""
|
||||
|
||||
def test_extracts_words(self):
|
||||
meta = _make_metadata(trained_words=["word1", "word2", "word3"])
|
||||
self.assertEqual(
|
||||
get_trigger_words_from_metadata(meta), ["word1", "word2", "word3"]
|
||||
)
|
||||
|
||||
def test_strips_whitespace(self):
|
||||
meta = _make_metadata(trained_words=[" padded ", "\ttabbed\t"])
|
||||
self.assertEqual(get_trigger_words_from_metadata(meta), ["padded", "tabbed"])
|
||||
|
||||
def test_filters_empty_strings(self):
|
||||
meta = _make_metadata(trained_words=["valid", "", " ", "also_valid"])
|
||||
self.assertEqual(get_trigger_words_from_metadata(meta), ["valid", "also_valid"])
|
||||
|
||||
def test_no_civitai_key(self):
|
||||
self.assertEqual(get_trigger_words_from_metadata({}), [])
|
||||
|
||||
def test_no_trained_words(self):
|
||||
meta = _make_metadata()
|
||||
self.assertEqual(get_trigger_words_from_metadata(meta), [])
|
||||
|
||||
def test_trained_words_not_list(self):
|
||||
meta = {"civitai": {"trainedWords": "not a list"}}
|
||||
self.assertEqual(get_trigger_words_from_metadata(meta), [])
|
||||
|
||||
def test_non_string_items_filtered(self):
|
||||
meta = _make_metadata(trained_words=["valid", 123, None, "also_valid"])
|
||||
self.assertEqual(get_trigger_words_from_metadata(meta), ["valid", "also_valid"])
|
||||
|
||||
|
||||
class TestGetExamplePrompt(unittest.TestCase):
|
||||
"""Test get_example_prompt_from_metadata — extracts first usable prompt."""
|
||||
|
||||
def test_extracts_first_prompt(self):
|
||||
images = [
|
||||
{"meta": {"prompt": "a beautiful landscape"}},
|
||||
{"meta": {"prompt": "second prompt"}},
|
||||
]
|
||||
meta = _make_metadata(images=images)
|
||||
self.assertEqual(
|
||||
get_example_prompt_from_metadata(meta), "a beautiful landscape"
|
||||
)
|
||||
|
||||
def test_skips_empty_prompts(self):
|
||||
images = [
|
||||
{"meta": {"prompt": ""}},
|
||||
{"meta": {"prompt": " "}},
|
||||
{"meta": {"prompt": "valid prompt"}},
|
||||
]
|
||||
meta = _make_metadata(images=images)
|
||||
self.assertEqual(get_example_prompt_from_metadata(meta), "valid prompt")
|
||||
|
||||
def test_no_images(self):
|
||||
meta = _make_metadata(images=[])
|
||||
self.assertIsNone(get_example_prompt_from_metadata(meta))
|
||||
|
||||
def test_no_civitai(self):
|
||||
self.assertIsNone(get_example_prompt_from_metadata({}))
|
||||
|
||||
def test_images_without_meta(self):
|
||||
images = [{"url": "http://example.com/img.jpg"}]
|
||||
meta = _make_metadata(images=images)
|
||||
self.assertIsNone(get_example_prompt_from_metadata(meta))
|
||||
|
||||
def test_meta_without_prompt(self):
|
||||
images = [{"meta": {"seed": 12345}}]
|
||||
meta = _make_metadata(images=images)
|
||||
self.assertIsNone(get_example_prompt_from_metadata(meta))
|
||||
|
||||
def test_non_dict_images_skipped(self):
|
||||
images = ["not a dict", None, {"meta": {"prompt": "found it"}}]
|
||||
meta = _make_metadata(images=images)
|
||||
self.assertEqual(get_example_prompt_from_metadata(meta), "found it")
|
||||
|
||||
def test_non_string_prompt_skipped(self):
|
||||
images = [{"meta": {"prompt": 12345}}, {"meta": {"prompt": "real prompt"}}]
|
||||
meta = _make_metadata(images=images)
|
||||
self.assertEqual(get_example_prompt_from_metadata(meta), "real prompt")
|
||||
|
||||
|
||||
class TestGetCivitaiImageUrls(unittest.TestCase):
|
||||
"""Test get_civitai_image_urls — extracts image URLs from metadata."""
|
||||
|
||||
def test_extracts_urls(self):
|
||||
images = [
|
||||
{"url": "https://civitai.com/img1.jpg"},
|
||||
{"url": "https://civitai.com/img2.jpg"},
|
||||
]
|
||||
meta = _make_metadata(images=images)
|
||||
urls = get_civitai_image_urls(meta)
|
||||
self.assertEqual(len(urls), 2)
|
||||
self.assertIn("https://civitai.com/img1.jpg", urls)
|
||||
|
||||
def test_filters_empty_urls(self):
|
||||
images = [{"url": ""}, {"url": "https://civitai.com/valid.jpg"}]
|
||||
meta = _make_metadata(images=images)
|
||||
urls = get_civitai_image_urls(meta)
|
||||
self.assertEqual(urls, ["https://civitai.com/valid.jpg"])
|
||||
|
||||
def test_no_images(self):
|
||||
meta = _make_metadata(images=[])
|
||||
self.assertEqual(get_civitai_image_urls(meta), [])
|
||||
|
||||
def test_no_civitai(self):
|
||||
self.assertEqual(get_civitai_image_urls({}), [])
|
||||
|
||||
def test_images_without_url_key(self):
|
||||
images = [{"id": 1}, {"url": "https://civitai.com/valid.jpg"}]
|
||||
meta = _make_metadata(images=images)
|
||||
urls = get_civitai_image_urls(meta)
|
||||
self.assertEqual(urls, ["https://civitai.com/valid.jpg"])
|
||||
|
||||
def test_non_dict_images_skipped(self):
|
||||
images = [None, "bad", {"url": "https://civitai.com/valid.jpg"}]
|
||||
meta = _make_metadata(images=images)
|
||||
urls = get_civitai_image_urls(meta)
|
||||
self.assertEqual(urls, ["https://civitai.com/valid.jpg"])
|
||||
|
||||
|
||||
# ── Filesystem-dependent tests ────────────────────────────────────────
|
||||
|
||||
|
||||
class TestReadLoraMetadata(unittest.TestCase):
|
||||
"""Test read_lora_metadata — file I/O with JSON parsing."""
|
||||
|
||||
def test_reads_valid_json(self):
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w", suffix=".metadata.json", delete=False
|
||||
) as f:
|
||||
json.dump({"civitai": {"trainedWords": ["test"]}}, f)
|
||||
f.flush()
|
||||
path = Path(f.name)
|
||||
try:
|
||||
result = read_lora_metadata(path)
|
||||
self.assertIsNotNone(result)
|
||||
self.assertEqual(result["civitai"]["trainedWords"], ["test"])
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_returns_none_for_invalid_json(self):
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w", suffix=".metadata.json", delete=False
|
||||
) as f:
|
||||
f.write("not valid json {{{")
|
||||
f.flush()
|
||||
path = Path(f.name)
|
||||
try:
|
||||
result = read_lora_metadata(path)
|
||||
self.assertIsNone(result)
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_returns_none_for_missing_file(self):
|
||||
result = read_lora_metadata(Path("/nonexistent/file.metadata.json"))
|
||||
self.assertIsNone(result)
|
||||
|
||||
|
||||
class TestGetLoraImageCacheDir(unittest.TestCase):
|
||||
"""Test get_lora_image_cache_dir — returns and creates cache path."""
|
||||
|
||||
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"))
|
||||
|
||||
def test_directory_exists(self):
|
||||
cache_dir = get_lora_image_cache_dir()
|
||||
self.assertTrue(cache_dir.is_dir())
|
||||
|
||||
|
||||
# ── TriggerWordCache tests ────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestTriggerWordCache(unittest.TestCase):
|
||||
"""Test TriggerWordCache — thread-safe trigger word lookup."""
|
||||
|
||||
def setUp(self):
|
||||
self.cache = TriggerWordCache()
|
||||
|
||||
def test_initial_state(self):
|
||||
self.assertFalse(self.cache.is_loaded)
|
||||
self.assertEqual(self.cache.get_trigger_words("anything"), [])
|
||||
|
||||
def test_load_from_temp_directory(self):
|
||||
"""Create temp metadata files and verify cache loads them."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# Create a metadata file
|
||||
meta = {
|
||||
"file_name": "my_lora.safetensors",
|
||||
"civitai": {"trainedWords": ["trigger1", "trigger2"]},
|
||||
}
|
||||
meta_path = Path(tmpdir) / "my_lora.safetensors.metadata.json"
|
||||
meta_path.write_text(json.dumps(meta))
|
||||
|
||||
# Patch find_lora_directories to return our temp dir
|
||||
with patch("py.lora_utils.find_lora_directories", return_value=[tmpdir]):
|
||||
count = self.cache.load(tmpdir)
|
||||
|
||||
self.assertTrue(self.cache.is_loaded)
|
||||
# Cache keys by both file_name stem and metadata filename stem
|
||||
self.assertGreaterEqual(count, 1)
|
||||
self.assertEqual(
|
||||
self.cache.get_trigger_words("my_lora"), ["trigger1", "trigger2"]
|
||||
)
|
||||
|
||||
def test_case_insensitive_lookup(self):
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
meta = {
|
||||
"file_name": "MyLoRA.safetensors",
|
||||
"civitai": {"trainedWords": ["word1"]},
|
||||
}
|
||||
(Path(tmpdir) / "MyLoRA.safetensors.metadata.json").write_text(
|
||||
json.dumps(meta)
|
||||
)
|
||||
|
||||
with patch("py.lora_utils.find_lora_directories", return_value=[tmpdir]):
|
||||
self.cache.load(tmpdir)
|
||||
|
||||
self.assertEqual(self.cache.get_trigger_words("mylora"), ["word1"])
|
||||
self.assertEqual(self.cache.get_trigger_words("MYLORA"), ["word1"])
|
||||
|
||||
def test_clear(self):
|
||||
# Manually set cache state
|
||||
self.cache._cache = {"test": ["word"]}
|
||||
self.cache._loaded = True
|
||||
|
||||
self.cache.clear()
|
||||
self.assertFalse(self.cache.is_loaded)
|
||||
self.assertEqual(self.cache.get_trigger_words("test"), [])
|
||||
|
||||
def test_unknown_lora_returns_empty(self):
|
||||
self.cache._cache = {"known": ["word"]}
|
||||
self.cache._loaded = True
|
||||
self.assertEqual(self.cache.get_trigger_words("unknown"), [])
|
||||
|
||||
def test_thread_safety(self):
|
||||
"""Verify concurrent access doesn't raise."""
|
||||
self.cache._cache = {"lora": ["word"]}
|
||||
self.cache._loaded = True
|
||||
|
||||
errors = []
|
||||
|
||||
def reader():
|
||||
try:
|
||||
for _ in range(100):
|
||||
self.cache.get_trigger_words("lora")
|
||||
except Exception as e:
|
||||
errors.append(e)
|
||||
|
||||
threads = [threading.Thread(target=reader) for _ in range(10)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
self.assertEqual(errors, [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user