198 lines
7.0 KiB
Python
198 lines
7.0 KiB
Python
import json
|
|
import sqlite3
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from core.contracts import SearchResult
|
|
from core.sources.lora_manager_archive import (
|
|
_build_result_from_row,
|
|
_extract_search_tokens,
|
|
_find_hash_match_row,
|
|
_load_full_rows_for_versions,
|
|
_normalize_model_type,
|
|
_query_candidate_rows,
|
|
search_lora_manager_archive_for_file,
|
|
)
|
|
|
|
|
|
class LoraManagerArchiveTests(unittest.TestCase):
|
|
|
|
def test_normalize_model_type(self):
|
|
self.assertEqual(_normalize_model_type("Loras"), "lora")
|
|
self.assertEqual(_normalize_model_type("Checkpoints"), "checkpoint")
|
|
self.assertEqual(_normalize_model_type("Embedding"), "textualinversion")
|
|
self.assertEqual(_normalize_model_type("Custom"), "custom")
|
|
|
|
def test_extract_search_tokens(self):
|
|
tokens = _extract_search_tokens("my_awesome_model_v1.safetensors")
|
|
self.assertIn("awesome", tokens["content_terms"])
|
|
self.assertIn("model", tokens["content_terms"])
|
|
self.assertIn("v1", tokens["version_terms"])
|
|
|
|
# Test that generic tokens and extensions are filtered out
|
|
tokens_generic = _extract_search_tokens("my_awesome_model_fp8_v1.sft")
|
|
self.assertNotIn("fp8", tokens_generic["content_terms"])
|
|
self.assertNotIn("sft", tokens_generic["content_terms"])
|
|
self.assertIn("awesome", tokens_generic["content_terms"])
|
|
self.assertIn("model", tokens_generic["content_terms"])
|
|
self.assertIn("v1", tokens_generic["version_terms"])
|
|
|
|
def test_query_candidate_rows_preserves_all_filter_plans(self):
|
|
class EmptyQueryResult:
|
|
def fetchall(self):
|
|
return []
|
|
|
|
class RecordingConnection:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def execute(self, sql, params):
|
|
self.calls.append((sql, tuple(params)))
|
|
return EmptyQueryResult()
|
|
|
|
connection = RecordingConnection()
|
|
|
|
candidates = _query_candidate_rows(
|
|
connection,
|
|
"aurora_style_v2.safetensors",
|
|
"lora",
|
|
2,
|
|
)
|
|
|
|
self.assertEqual([], candidates)
|
|
self.assertEqual(5, len(connection.calls))
|
|
self.assertTrue(
|
|
all("m.type = ?" in sql for sql, _params in connection.calls)
|
|
)
|
|
self.assertTrue(
|
|
all(
|
|
"ORDER BY m.id ASC, COALESCE(v.position, 999999) ASC LIMIT ?"
|
|
in sql
|
|
for sql, _params in connection.calls
|
|
)
|
|
)
|
|
self.assertEqual(
|
|
[
|
|
("%aurora%", "%style%", "%v2%", "LORA", 50),
|
|
("%aurora%", "%style%", "LORA", 50),
|
|
("%aurora%", "%v2%", "LORA", 50),
|
|
("%aurora%", "LORA", 50),
|
|
("%v2%", "LORA", 50),
|
|
],
|
|
[params for _sql, params in connection.calls],
|
|
)
|
|
|
|
def test_database_operations_with_in_memory_sqlite(self):
|
|
# Create an in-memory database
|
|
conn = sqlite3.connect(":memory:")
|
|
conn.row_factory = sqlite3.Row
|
|
|
|
# Create tables
|
|
conn.execute("""
|
|
CREATE TABLE models (
|
|
id INTEGER PRIMARY KEY,
|
|
name TEXT,
|
|
type TEXT,
|
|
data TEXT
|
|
)
|
|
""")
|
|
conn.execute("""
|
|
CREATE TABLE model_versions (
|
|
id INTEGER PRIMARY KEY,
|
|
model_id INTEGER,
|
|
name TEXT,
|
|
base_model TEXT,
|
|
position INTEGER,
|
|
data TEXT
|
|
)
|
|
""")
|
|
conn.execute("""
|
|
CREATE TABLE model_files (
|
|
id INTEGER PRIMARY KEY,
|
|
version_id INTEGER,
|
|
type TEXT,
|
|
data TEXT
|
|
)
|
|
""")
|
|
|
|
# Insert dummy data
|
|
conn.execute("INSERT INTO models VALUES (1, 'ghost_checkpoint', 'CHECKPOINT', ?)", (json.dumps({"tags": ["anime"]}),))
|
|
conn.execute("INSERT INTO model_versions VALUES (10, 1, 'v1.0', 'SD 1.5', 1, ?)", (json.dumps({"trainedWords": ["ghost_word"]}),))
|
|
archive_hash = "a" * 64
|
|
conn.execute(
|
|
"INSERT INTO model_files VALUES (100, 10, 'Model', ?)",
|
|
(
|
|
json.dumps(
|
|
{
|
|
"name": "ghost_v1.safetensors",
|
|
"downloadUrl": "https://civitai.com/api/download/models/10",
|
|
"sizeKB": 3,
|
|
"hashes": {"SHA256": archive_hash},
|
|
}
|
|
),
|
|
),
|
|
)
|
|
conn.commit()
|
|
|
|
# Test querying candidates
|
|
candidates = _query_candidate_rows(conn, "ghost_v1.safetensors", "checkpoint", 10)
|
|
self.assertEqual(len(candidates), 1)
|
|
self.assertEqual(candidates[0]["version_id"], 10)
|
|
|
|
# Test loading full version rows
|
|
full_rows = _load_full_rows_for_versions(conn, [10])
|
|
self.assertIn(10, full_rows)
|
|
row = full_rows[10]
|
|
self.assertEqual(row["model_name"], "ghost_checkpoint")
|
|
self.assertEqual(row["version_name"], "v1.0")
|
|
|
|
# Test building result from row
|
|
result = _build_result_from_row(conn, row, "ghost_checkpoint")
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result.filename, "ghost_v1.safetensors")
|
|
self.assertEqual(result.name, "ghost_checkpoint")
|
|
self.assertEqual(list(result.trained_words), ["ghost_word"])
|
|
self.assertEqual(result.size, 3 * 1024)
|
|
self.assertEqual(result.sha256, archive_hash)
|
|
self.assertEqual(dict(result.hashes), {"SHA256": archive_hash})
|
|
|
|
hash_match = _find_hash_match_row(conn, archive_hash)
|
|
self.assertIsNotNone(hash_match)
|
|
self.assertEqual(hash_match[0]["version_id"], 10)
|
|
self.assertEqual(hash_match[1]["name"], "ghost_v1.safetensors")
|
|
|
|
with patch(
|
|
"core.sources.lora_manager_archive._connect_readonly",
|
|
return_value=conn,
|
|
):
|
|
hash_result = search_lora_manager_archive_for_file(
|
|
"",
|
|
sha256=archive_hash,
|
|
)
|
|
|
|
self.assertIsNotNone(hash_result)
|
|
self.assertEqual(hash_result.match_type, "hash")
|
|
self.assertEqual(hash_result.confidence, 100.0)
|
|
self.assertEqual(hash_result.sha256, archive_hash)
|
|
|
|
def test_filename_search_keeps_typed_candidates(self):
|
|
filename = "typed-lora-archive-model.safetensors"
|
|
candidate = SearchResult(
|
|
source="lora_manager_archive",
|
|
name="typed-lora-archive-model",
|
|
version_name="v1",
|
|
filename=filename,
|
|
download_url="https://example.com/typed-lora-archive-model.safetensors",
|
|
)
|
|
|
|
with patch(
|
|
"core.sources.lora_manager_archive.search_lora_manager_archive",
|
|
return_value=[candidate],
|
|
):
|
|
result = search_lora_manager_archive_for_file(filename)
|
|
|
|
self.assertIsInstance(result, SearchResult)
|
|
self.assertEqual(filename, result.filename)
|
|
self.assertEqual(100.0, result.confidence)
|
|
self.assertEqual("exact", result.match_type)
|