import unittest from unittest.mock import patch from core.contracts import SearchResult from core.matcher import ( build_filename_search_queries, calculate_archived_model_confidence, calculate_filename_confidence, calculate_similarity, find_matches, normalize_filename, normalize_model_family_filename, ) from core.progress import report_progress from core.sources.model_list import search_model_list, search_model_list_multiple from core.sources.popular import search_popular_models from core.type_utils import ( MODEL_CONTAINER_SUFFIXES, MODEL_VARIANT_SUFFIXES, NUMERIC_PRECISION_SUFFIXES, QUANTIZATION_LEVEL_SUFFIXES, QUANTIZATION_SCHEME_SUFFIXES, TECHNICAL_MODEL_SUFFIXES, as_dict, as_list, format_size_bytes, ) class MatcherTests(unittest.TestCase): def test_normalize_filename_removes_extension_and_converts_lowercase(self): self.assertEqual("test model", normalize_filename("Test_Model.safetensors")) self.assertEqual("test model", normalize_filename("test-model.ckpt")) def test_normalize_filename_replaces_dots_and_separators(self): self.assertEqual("test model v1 0", normalize_filename("test.model.v1.0.safetensors")) self.assertEqual("my cool model 1 2 3", normalize_filename("My_Cool-Model.1.2.3.ckpt")) def test_calculate_similarity_identical(self): self.assertEqual(1.0, calculate_similarity("hello", "hello")) def test_calculate_similarity_different(self): self.assertLess(calculate_similarity("hello", "world"), 0.5) def test_find_matches_exact_and_fuzzy(self): candidates = [ {"filename": "sd_xl_base_1.0.safetensors"}, {"filename": "pony Diffusion.safetensors"}, ] matches = find_matches("sd_xl_base_1.0.safetensors", candidates) self.assertEqual(2, len(matches)) self.assertEqual("sd_xl_base_1.0.safetensors", matches[0].filename) self.assertEqual(1.0, matches[0].similarity) # With higher threshold, only the best match should be returned matches_filtered = find_matches("sd_xl_base_1.0.safetensors", candidates, threshold=0.7) self.assertEqual(1, len(matches_filtered)) def test_model_family_normalization_ignores_precision_and_quantization(self): self.assertEqual( "qwen3vl 4b abliterated", normalize_model_family_filename( "qwen3vl-4b-abliterated_fp8_e4m3fn_scaled.safetensors" ), ) self.assertEqual( "qwen3vl 4b abliterated", normalize_model_family_filename( "qwen3vl-4b-abliterated_int8_conv.safetensors" ), ) self.assertEqual( "qwen3vl 4b", normalize_model_family_filename("qwen3vl_4b_iq4_nl.gguf"), ) def test_model_family_normalization_handles_common_technical_families(self): suffixes = [ "float16", "half", "e4m3fnuz", "e5m2fnuz", "e2m1", "nf4", "q4_0", "q3_k_m", "iq2_xxs", "iq4_nl", "tq1_0", "mxfp4", "nvfp4", "gptq", "awq", "exl2", "hqq", "aqlm", "bnb4", "ema-only", "pruned", "convrot", ] for suffix in suffixes: with self.subTest(suffix=suffix): self.assertEqual( "qwen3vl 4b", normalize_model_family_filename( f"qwen3vl_4b_{suffix}.safetensors" ), ) def test_model_family_normalization_handles_compact_precision_variants(self): self.assertIn( "Krea2_Turbo", build_filename_search_queries( "Krea2_Turbo_convrot_int8mixed.safetensors" ), ) self.assertEqual( "krea2 turbo", normalize_model_family_filename( "Krea2_Turbo_convrot_int8mixed.safetensors" ), ) self.assertEqual( "krea2 raw", normalize_model_family_filename( "Krea2_Raw_convrot_int8mixed.safetensors" ), ) def test_find_matches_includes_krea2_conversion_variants(self): target = "krea2_turbo_bf16.safetensors" candidates = [ {"filename": "Krea2_Raw_convrot_int8mixed.safetensors"}, {"filename": "Krea2_Turbo_convrot_int8mixed.safetensors"}, ] matches = find_matches(target, candidates, threshold=0.7) self.assertEqual( [ "Krea2_Turbo_convrot_int8mixed.safetensors", "Krea2_Raw_convrot_int8mixed.safetensors", ], [match.filename for match in matches], ) self.assertEqual(94.0, matches[0].confidence) self.assertEqual(70.0, matches[1].confidence) def test_filename_confidence_prioritizes_family_and_size_over_precision(self): target = "qwen3vl-4b-abliterated_fp8_e4m3fn.safetensors" self.assertEqual( 94.0, calculate_filename_confidence( target, "qwen3vl-4b-abliterated_int8.safetensors" ), ) self.assertGreaterEqual( calculate_filename_confidence(target, "qwen3vl_4b_int8.safetensors"), 90.0, ) self.assertGreaterEqual( calculate_filename_confidence( target, "qwen3vl-4b-heretic_int8.safetensors" ), 84.0, ) self.assertLess( calculate_filename_confidence( target, "qwen_image_fp8_e4m3fn.safetensors" ), calculate_filename_confidence( target, "qwen3vl-4b-heretic_int8.safetensors" ), ) def test_filename_confidence_ignores_token_order(self): target = "yue2_convrot_int8.safetensors" confidence = calculate_filename_confidence( target, "yue2_3b_int8_convrot.safetensors" ) self.assertGreaterEqual(confidence, 90.0) def test_filename_confidence_does_not_merge_different_sizes_or_generations(self): target = "qwen3vl-4b_fp8_e4m3fn.safetensors" self.assertLess( calculate_filename_confidence(target, "qwen3vl-8b_int8.safetensors"), 70.0, ) self.assertLess( calculate_filename_confidence(target, "qwen2vl-4b_int8.safetensors"), 70.0, ) def test_find_matches_includes_compatible_qwen_precision_variants(self): target = "qwen3vl-4b-abliterated_fp8_e4m3fn.safetensors" candidates = [ {"filename": "qwen3vl-4b-abliterated_int8.safetensors"}, {"filename": "qwen3vl_4b_int8.safetensors"}, {"filename": "qwen3vl-4b-heretic_int8.safetensors"}, {"filename": "qwen3vl-8b_int8.safetensors"}, {"filename": "qwen_image_fp8_e4m3fn.safetensors"}, ] matches = find_matches(target, candidates, threshold=0.7) self.assertEqual( [ "qwen3vl-4b-abliterated_int8.safetensors", "qwen3vl_4b_int8.safetensors", "qwen3vl-4b-heretic_int8.safetensors", "qwen_image_fp8_e4m3fn.safetensors", ], [match.filename for match in matches], ) def test_find_matches_does_not_mutate_candidate_mappings(self): candidate = {"filename": "model.safetensors", "category": "checkpoints"} matches = find_matches("model.safetensors", [candidate]) self.assertEqual(1, len(matches)) self.assertEqual( {"filename": "model.safetensors", "category": "checkpoints"}, candidate, ) self.assertEqual("checkpoints", matches[0].model.category) def test_search_model_list_unification(self): try: search_model_list("nonexistent_model_name_xyz.safetensors", exact_only=True) search_model_list_multiple("nonexistent_model_name_xyz.safetensors") except Exception as e: self.fail(f"search_model_list failed with exception: {e}") def test_model_list_result_fields_preserve_match_mode_contracts(self): models = [ { "filename": "exact.safetensors", "url": "https://example.test/exact", "name": "Exact model", "type": "checkpoints", "save_path": "checkpoints", "size": "12 MB", }, { "filename": "target-model.safetensors", "url": "https://example.test/target", "name": "Target model", "type": "loras", "save_path": "loras", "size": 24, }, ] with patch("core.sources.model_list._load_model_list", return_value=models): exact = search_model_list("exact.safetensors", exact_only=True) fuzzy = search_model_list("target-model-extra.safetensors") multiple = search_model_list_multiple("exact", limit=1) self.assertIsInstance(exact, SearchResult) self.assertEqual(exact.source, "model_list") self.assertEqual(exact.filename, "exact.safetensors") self.assertEqual(exact.url, "https://example.test/exact") self.assertEqual(exact.name, "Exact model") self.assertEqual(exact.model_type, "checkpoints") self.assertEqual(exact.extra_value("directory"), "checkpoints") self.assertEqual(exact.size, 12 * 1024**2) self.assertEqual(exact.match_type, "exact") self.assertEqual(fuzzy.filename, "target-model.safetensors") self.assertEqual(fuzzy.match_type, "fuzzy") self.assertGreater(fuzzy.confidence, 0) self.assertEqual(multiple[0].filename, "exact.safetensors") self.assertGreater(multiple[0].confidence, 0) self.assertEqual(multiple[0].match_type, "similar") def test_popular_search_returns_typed_search_results(self): with patch( "core.sources.popular._load_popular_models", return_value={ "popular-model.safetensors": { "url": "https://example.test/popular", "type": "checkpoints", "directory": "checkpoints", "canonical_name": "popular-model", } }, ): results = search_popular_models("popular-model", limit=1) self.assertEqual(1, len(results)) self.assertIsInstance(results[0], SearchResult) self.assertEqual("popular-model.safetensors", results[0].filename) self.assertEqual("https://example.test/popular", results[0].url) self.assertEqual("checkpoints", results[0].model_type) self.assertEqual("checkpoints", results[0].extra_value("directory")) self.assertEqual("popular-model", results[0].extra_value("canonical_name")) def test_build_filename_search_queries(self): # Verify that common precisions are extracted correctly queries = build_filename_search_queries("sdxl_ep10_fp16.safetensors") self.assertIn("sdxl_ep10_fp16.safetensors", queries) self.assertIn("sdxl_ep10_fp16", queries) self.assertIn("sdxl_ep10", queries) queries_simple = build_filename_search_queries("model.ckpt") self.assertIn("model.ckpt", queries_simple) self.assertIn("model", queries_simple) queries_int8 = build_filename_search_queries( "qwen3vl_4b_abliterated_int8.safetensors" ) self.assertIn("qwen3vl_4b_abliterated", queries_int8) for filename in [ "qwen3vl_4b_q4_k_m.gguf", "qwen3vl_4b_iq2_xxs.gguf", "qwen3vl_4b_e4m3fnuz.safetensors", "qwen3vl_4b_exl2.safetensors", ]: with self.subTest(filename=filename): self.assertIn("qwen3vl_4b", build_filename_search_queries(filename)) def test_archived_model_confidence_uses_model_family_variants(self): target = "qwen3vl-4b-abliterated_fp8_e4m3fn.safetensors" self.assertEqual( 94.0, calculate_archived_model_confidence( target, filename="qwen3vl-4b-abliterated_int8.safetensors", ), ) self.assertGreaterEqual( calculate_archived_model_confidence( target, filename="qwen3vl-4b-heretic_int8.safetensors", ), 84.0, ) self.assertLess( calculate_archived_model_confidence( target, filename="qwen3vl-8b_int8.safetensors", ), 70.0, ) class ProgressTests(unittest.TestCase): def test_report_progress_calls_callback(self): called_payload = None def callback(payload): nonlocal called_payload called_payload = payload report_progress( callback, stage="test_stage", message="test_message", percent=50.0, foo="bar" ) self.assertIsNotNone(called_payload) self.assertEqual("test_stage", called_payload.get("stage")) self.assertEqual("test_message", called_payload.get("message")) self.assertEqual(50.0, called_payload.get("percent")) self.assertEqual("bar", called_payload.get("foo")) def test_report_progress_handles_exceptions(self): def bad_callback(payload): raise ValueError("Boom") # Should not raise exception try: report_progress(bad_callback, "stage", "msg") except Exception as e: self.fail(f"report_progress raised exception: {e}") def test_report_progress_handles_none_callback(self): # Should do nothing and not raise exception try: report_progress(None, "stage", "msg") except Exception as e: self.fail(f"report_progress raised exception with None callback: {e}") class TypeUtilsTests(unittest.TestCase): def test_technical_suffix_categories_are_disjoint_and_combined(self): precision = set(NUMERIC_PRECISION_SUFFIXES) quantization_levels = set(QUANTIZATION_LEVEL_SUFFIXES) quantization_schemes = set(QUANTIZATION_SCHEME_SUFFIXES) containers = set(MODEL_CONTAINER_SUFFIXES) variants = set(MODEL_VARIANT_SUFFIXES) categories = [ precision, quantization_levels, quantization_schemes, containers, variants, ] for index, category in enumerate(categories): for other_category in categories[index + 1 :]: self.assertTrue(category.isdisjoint(other_category)) self.assertEqual( set().union(*categories), set(TECHNICAL_MODEL_SUFFIXES), ) def test_as_dict_returns_dict_when_dict(self): d = {"a": 1} self.assertEqual(d, as_dict(d)) def test_as_dict_returns_empty_dict_when_not_dict(self): self.assertEqual({}, as_dict("not a dict")) self.assertEqual({}, as_dict(None)) def test_as_list_returns_filtered_list(self): self.assertEqual([1, 2], as_list([1, None, "", 2])) def test_as_list_handles_tuples_and_sets(self): self.assertEqual([1, 2], as_list((1, None, "", 2))) self.assertEqual([1, 2], as_list({1, "", 2})) def test_as_list_handles_comma_separated_strings(self): self.assertEqual(["a", "b", "c"], as_list("a, b,, c")) def test_as_list_returns_empty_list_otherwise(self): self.assertEqual([], as_list(None)) self.assertEqual([], as_list(123)) def test_format_size_bytes(self): self.assertEqual("1.5 GB", format_size_bytes(1.5 * 1024 * 1024 * 1024)) self.assertEqual("1.5GB", format_size_bytes(1.5 * 1024 * 1024 * 1024, include_space=False)) self.assertEqual("0 B", format_size_bytes(0)) self.assertEqual("0B", format_size_bytes(0, include_space=False)) self.assertIsNone(format_size_bytes(None)) self.assertIsNone(format_size_bytes("")) self.assertEqual("invalid_str", format_size_bytes("invalid_str")) self.assertEqual("500.0 KB", format_size_bytes("512000")) if __name__ == "__main__": unittest.main()