42 lines
1.5 KiB
Python
42 lines
1.5 KiB
Python
import unittest
|
|
|
|
from core.type_utils import normalize_download_category, resolve_model_category
|
|
|
|
|
|
class CategoryMappingTests(unittest.TestCase):
|
|
def test_folder_category_aliases_resolve_to_existing_comfyui_keys(self):
|
|
expected = {
|
|
"checkpoint": "checkpoints",
|
|
"lora": "loras",
|
|
"embedding": "embeddings",
|
|
"textual inversion": "embeddings",
|
|
"textualinversion": "embeddings",
|
|
"upscaler": "upscale_models",
|
|
"unet gguf": "diffusion_models",
|
|
"select-safetensors": "diffusion_models",
|
|
"clip": "text_encoders",
|
|
"ip_adapter": "ipadapter",
|
|
"ultralytics_bbox": "ultralytics",
|
|
"background-removal-model": "background_removal",
|
|
"default": "upscale_models",
|
|
}
|
|
|
|
for raw_category, canonical_category in expected.items():
|
|
self.assertEqual(
|
|
canonical_category,
|
|
resolve_model_category(raw_category, target_format="folder"),
|
|
)
|
|
self.assertEqual(
|
|
canonical_category,
|
|
normalize_download_category(raw_category),
|
|
)
|
|
|
|
def test_unknown_and_empty_categories_keep_current_fallbacks(self):
|
|
self.assertEqual("checkpoints", normalize_download_category(""))
|
|
self.assertEqual("checkpoints", normalize_download_category(None))
|
|
self.assertEqual("new_category", normalize_download_category("new category"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|