317 lines
11 KiB
Python
317 lines
11 KiB
Python
"""
|
|
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 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()
|