107 lines
4.0 KiB
Python
107 lines
4.0 KiB
Python
import unittest
|
|
import sqlite3
|
|
import json
|
|
from core.sources.lora_manager_archive import (
|
|
_normalize_model_type,
|
|
_extract_search_tokens,
|
|
_query_candidate_rows,
|
|
_load_full_rows_for_versions,
|
|
_build_result_from_row,
|
|
)
|
|
|
|
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_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",
|
|
"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(result["trained_words"], ["ghost_word"])
|
|
self.assertEqual(result["sha256"], archive_hash)
|
|
self.assertEqual(result["hashes"], {"SHA256": archive_hash})
|
|
|
|
conn.close()
|