Azornes-Comfyui-Model-Resolver/tests/test_huggingface.py

import unittest
import re
import json
import tempfile
from pathlib import Path
from unittest.mock import patch
from core.sources.huggingface import (
    _fetch_author_index,
    _is_author_index_fresh,
    _read_persistent_author_indexes,
    build_huggingface_custom_result,
    parse_huggingface_url,
    get_huggingface_download_url,
    _normalize_huggingface_size_probe_url,
    _write_persistent_author_index,
)
from core.matcher import build_filename_search_queries
from core.type_utils import extract_file_size

class HuggingFaceSourceTests(unittest.TestCase):

    def test_fetch_author_index_builds_index_from_provider_repo_list(self):
        with patch(
            "core.sources.huggingface.execute_provider_json_request",
            return_value=[
                {
                    "id": "Comfy-Org/example",
                    "siblings": [
                        {
                            "rfilename": "text_encoders/model.safetensors",
                            "size": 123,
                        }
                    ],
                }
            ],
        ):
            index = _fetch_author_index("Comfy-Org", headers={})

        self.assertIsInstance(index, dict)
        self.assertEqual("Comfy-Org", index["author"])
        self.assertEqual(1, index["repo_count"])
        self.assertEqual(1, index["file_count"])
        self.assertEqual(
            "text_encoders/model.safetensors",
            index["files"][0]["path"],
        )

    def test_malformed_list_author_index_is_not_treated_as_fresh(self):
        self.assertFalse(_is_author_index_fresh([{"id": "Comfy-Org/example"}]))

    def test_persistent_cache_discards_raw_provider_list_entries(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            index_path = Path(tmpdir) / "huggingface-author-index.json"
            index_path.write_text(
                json.dumps(
                    {
                        "version": 1,
                        "authors": {
                            "Comfy-Org": [{"id": "Comfy-Org/example"}],
                            "Valid": {
                                "author": "Valid",
                                "updated_at": 1,
                                "repo_count": 0,
                                "file_count": 0,
                                "repos": [],
                                "files": [],
                            },
                        },
                    }
                ),
                encoding="utf-8",
            )
            with patch(
                "core.sources.huggingface.HF_AUTHOR_INDEX_CACHE_PATH",
                str(index_path),
            ):
                data = _read_persistent_author_indexes()

        self.assertNotIn("Comfy-Org", data["authors"])
        self.assertIn("Valid", data["authors"])

    def test_persistent_author_index_is_pretty_printed(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            index_path = Path(tmpdir) / "huggingface-author-index.json"
            with patch(
                "core.sources.huggingface.HF_AUTHOR_INDEX_CACHE_PATH",
                str(index_path),
            ):
                _write_persistent_author_index(
                    "Example",
                    {
                        "author": "Example",
                        "updated_at": 1,
                        "repo_count": 0,
                        "file_count": 0,
                        "repos": [],
                        "files": [],
                    },
                )

            content = index_path.read_text(encoding="utf-8")
            self.assertIn('\n  "version":', content)
            self.assertIn('\n  "authors": {', content)
            self.assertTrue(content.endswith("\n"))
            self.assertEqual("Example", json.loads(content)["authors"]["Example"]["author"])

    def test_parse_huggingface_url_valid_http(self):
        url = "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors"
        result = parse_huggingface_url(url)
        self.assertIsNotNone(result)
        self.assertEqual(result["repo"], "runwayml/stable-diffusion-v1-5")
        self.assertEqual(result["branch"], "main")
        self.assertEqual(result["filename"], "v1-5-pruned-emaonly.safetensors")

    def test_parse_huggingface_url_valid_hf_protocol(self):
        url = "hf://stabilityai/stable-diffusion-xl-base-1.0/sd_xl_base_1.0.safetensors"
        result = parse_huggingface_url(url)
        self.assertIsNotNone(result)
        self.assertEqual(result["repo"], "stabilityai/stable-diffusion-xl-base-1.0")
        self.assertEqual(result["filename"], "sd_xl_base_1.0.safetensors")

    def test_parse_huggingface_url_invalid(self):
        self.assertIsNone(parse_huggingface_url("https://example.com/not-hf"))
        self.assertIsNone(parse_huggingface_url("hf://not_enough_slashes"))

    def test_get_huggingface_download_url(self):
        repo = "runwayml/stable-diffusion-v1-5"
        filename = "v1-5-pruned-emaonly.safetensors"
        url = get_huggingface_download_url(repo, filename)
        self.assertEqual(
            url,
            "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors"
        )

    def test_extract_file_size(self):
        # Directly nested key
        self.assertEqual(extract_file_size({"size": 100}), 100)
        self.assertEqual(extract_file_size({"sizeBytes": "200"}), 200)
        
        # LFS nested key
        self.assertEqual(extract_file_size({"lfs": {"size": 300}}), 300)
        
        # None cases
        self.assertIsNone(extract_file_size(None))
        self.assertIsNone(extract_file_size({}))

    def test_normalize_huggingface_size_probe_url(self):
        url = "https://huggingface.co/user/repo/blob/main/model.safetensors"
        normalized = _normalize_huggingface_size_probe_url(url)
        self.assertEqual(normalized, "https://huggingface.co/user/repo/resolve/main/model.safetensors")

    def test_custom_result_uses_sha256_from_exact_huggingface_file(self):
        sha256 = "a" * 64
        url = (
            "https://huggingface.co/DreamFast/Qwen3-VL-4b-Heretic-ComfyUI/"
            "blob/main/qwen3-vl-4b-heretic_int8.safetensors"
        )
        with patch(
            "core.sources.huggingface._get_repo_tree",
            return_value=[
                {
                    "path": "qwen3-vl-4b-heretic_int8.safetensors",
                    "size": 123,
                    "lfs": {"oid": f"sha256:{sha256}", "size": 123},
                }
            ],
        ):
            result = build_huggingface_custom_result(url)

        self.assertIsNotNone(result)
        self.assertEqual(sha256, result["sha256"])
        self.assertEqual({"SHA256": sha256}, result["hashes"])
        self.assertEqual(123, result["size"])
        self.assertEqual("huggingface", result["source"])
        self.assertTrue(result["custom_url"])

    def test_custom_result_leaves_hash_empty_when_huggingface_does_not_provide_it(self):
        url = "https://huggingface.co/user/repo/blob/main/model.safetensors"
        with patch("core.sources.huggingface._get_repo_tree", return_value=None):
            with patch(
                "core.sources.huggingface.fetch_remote_file_size_cached",
                return_value=456,
            ):
                result = build_huggingface_custom_result(url)

        self.assertIsNotNone(result)
        self.assertEqual("", result["sha256"])
        self.assertEqual({}, result["hashes"])
        self.assertEqual(456, result["size"])

    def test_build_huggingface_search_queries(self):
        queries = build_filename_search_queries("some_model_bf16.safetensors")
        self.assertIn("some_model", queries)
        self.assertIn("some_model_bf16", queries)