53 lines
2.0 KiB
Python
53 lines
2.0 KiB
Python
import json
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from core.type_utils import CATEGORY_MAP, 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"))
|
|
|
|
def test_frontend_category_alias_artifact_matches_backend_map(self):
|
|
project_root = Path(__file__).resolve().parents[1]
|
|
artifact = (
|
|
project_root / "web" / "resolver" / "utils" / "category_aliases.generated.js"
|
|
).read_text(encoding="utf-8")
|
|
payload = artifact.split("Object.freeze(", 1)[1].rsplit(");", 1)[0].strip()
|
|
|
|
self.assertEqual(CATEGORY_MAP, json.loads(payload))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|