92 lines
3.8 KiB
Python
92 lines
3.8 KiB
Python
import unittest
|
|
import re
|
|
import json
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
from core.sources.huggingface import (
|
|
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_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_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)
|