from __future__ import annotations import sys import unittest from pathlib import Path from unittest.mock import patch PACKAGE_DIR = Path(__file__).resolve().parents[1] / "modelhub_submmit_api" if str(PACKAGE_DIR) in sys.path: sys.path.remove(str(PACKAGE_DIR)) sys.path.insert(0, str(PACKAGE_DIR)) from hf_discovery import HuggingFaceDiscovery, inspect_repo_tree # noqa: E402 from http_json import HttpJsonError # noqa: E402 class RecordingHttpClient: def __init__(self) -> None: self.queries: list[dict] = [] def request_json(self, _method: str, _path: str, *, query: dict) -> dict: self.queries.append(query) return {"success": True, "data": {"models": []}} class RateLimitedPageClient: def __init__(self) -> None: self.page_calls: list[int] = [] self.page_two_attempts = 0 def request_json(self, _method: str, _path: str, *, query: dict) -> dict: page = int(query["page_number"]) self.page_calls.append(page) if page == 2: self.page_two_attempts += 1 if self.page_two_attempts == 1: raise HttpJsonError("rate limited", status_code=429) start = (page - 1) * 50 items = [ { "id": f"owner/model-{index}", "downloads": 100, "last_modified": "2026-01-01T00:00:00Z", } for index in range(start, start + 50) ] return {"success": True, "data": {"models": items}} class ModelMetadataClient: def __init__(self) -> None: self.calls = 0 def request_json(self, _method: str, _path: str, **_kwargs) -> dict: # noqa: ANN003 self.calls += 1 return {"Code": 200, "Data": {"LastUpdatedTime": 1786294389}} class ModelScopeDiscoveryTests(unittest.TestCase): def test_listing_diagnostics_distinguish_filtered_prefix_from_exhaustion(self) -> None: class FilteredPageClient: def request_json(self, *_args, **_kwargs): return {"data": {"models": [ {"id": f"owner/model-{index}", "downloads": 0} for index in range(50) ]}} discovery = HuggingFaceDiscovery(http_client=FilteredPageClient()) # type: ignore[arg-type] self.assertEqual([], discovery.list_recent_models( pipeline_tags=["text-generation"], limit=50, min_downloads=50, )) diagnostic = discovery.last_listing_diagnostics[0] self.assertEqual(50, diagnostic["rawRowsScanned"]) self.assertEqual(0, diagnostic["acceptedModels"]) self.assertEqual("page_budget", diagnostic["stopReason"]) self.assertFalse(diagnostic["catalogExhausted"]) def test_failed_repo_tree_is_retried_instead_of_cached_as_empty(self) -> None: class RecoveringClient: calls = 0 def request_json(self, *_args, **_kwargs): self.calls += 1 if self.calls == 1: raise HttpJsonError("temporary outage", status_code=503) return {"Data": {"Files": [{"Path": "model.safetensors", "Size": 1024}]}} client = RecoveringClient() discovery = HuggingFaceDiscovery(legacy_http_client=client) # type: ignore[arg-type] with self.assertRaises(HttpJsonError): discovery.list_repo_tree("owner/recoverable") self.assertNotIn("owner/recoverable", discovery._repo_tree_cache) entries = discovery.list_repo_tree("owner/recoverable") self.assertEqual("model.safetensors", entries[0]["Path"]) self.assertEqual(entries, discovery.list_repo_tree("owner/recoverable")) self.assertEqual(2, client.calls) def test_duplicate_normalized_paths_do_not_inflate_repository_size(self) -> None: gib = 1024 ** 3 inspection = inspect_repo_tree("owner/model", [ {"Path": "model.safetensors", "Size": 20 * gib}, {"Path": "./model.safetensors", "Size": 19 * gib}, {"Path": "/model.safetensors", "Size": 20 * gib}, ]) self.assertEqual(["model.safetensors"], inspection.file_paths) self.assertEqual(["model.safetensors"], inspection.weight_files) self.assertEqual(20 * gib, inspection.repository_size_bytes) self.assertEqual(20 * gib, inspection.estimated_load_bytes("vllm")) def test_model_last_modified_uses_detail_api_and_is_cached(self) -> None: metadata_client = ModelMetadataClient() discovery = HuggingFaceDiscovery(legacy_http_client=metadata_client) # type: ignore[arg-type] first = discovery.get_model_last_modified("owner/model") second = discovery.get_model_last_modified("owner/model") self.assertEqual(first, second) self.assertEqual(1786294389, int(first.timestamp())) # type: ignore[union-attr] self.assertEqual(1, metadata_client.calls) def test_model_detail_cache_evicts_old_entries_at_fixed_limit(self) -> None: metadata_client = ModelMetadataClient() with patch.dict("os.environ", {"MODELSCOPE_DETAIL_CACHE_MAX_MODELS": "32"}): discovery = HuggingFaceDiscovery(legacy_http_client=metadata_client) # type: ignore[arg-type] for index in range(40): discovery.get_model_last_modified(f"owner/model-{index}") self.assertEqual(32, len(discovery._model_last_modified_cache)) discovery.get_model_last_modified("owner/model-0") self.assertEqual(41, metadata_client.calls) def test_openapi_page_size_never_exceeds_platform_limit(self) -> None: http_client = RecordingHttpClient() discovery = HuggingFaceDiscovery(http_client=http_client) # type: ignore[arg-type] models = discovery.list_recent_models( pipeline_tags=["text-generation"], limit=1000, min_downloads=0, ) self.assertEqual([], models) self.assertEqual(50, http_client.queries[0]["page_size"]) def test_popular_sort_uses_separate_page_cache(self) -> None: http_client = RecordingHttpClient() discovery = HuggingFaceDiscovery(http_client=http_client) # type: ignore[arg-type] discovery.list_recent_models(pipeline_tags=["text-generation"], limit=50, min_downloads=0) discovery.list_recent_models( pipeline_tags=["text-generation"], limit=50, min_downloads=0, sort="downloads" ) self.assertEqual(["last_modified", "downloads"], [item["sort"] for item in http_client.queries]) def test_popular_scan_rotates_bounded_page_windows(self) -> None: http_client = RecordingHttpClient() discovery = HuggingFaceDiscovery(http_client=http_client) # type: ignore[arg-type] for _ in range(2): discovery.list_recent_models( pipeline_tags=["text-generation"], limit=200, min_downloads=50, sort="downloads", rotate_pages=True, ) self.assertEqual([1, 5], [item["page_number"] for item in http_client.queries]) def test_rate_limited_page_keeps_previous_pages_and_resumes_from_cache(self) -> None: http_client = RateLimitedPageClient() discovery = HuggingFaceDiscovery( http_client=http_client, # type: ignore[arg-type] page_interval_seconds=0, page_cache_ttl_seconds=60, ) partial = discovery.list_recent_models( pipeline_tags=["text-generation"], limit=100, min_downloads=0, ) resumed = discovery.list_recent_models( pipeline_tags=["text-generation"], limit=100, min_downloads=0, ) self.assertEqual(50, len(partial)) self.assertEqual(100, len(resumed)) self.assertEqual([1, 2, 2], http_client.page_calls) if __name__ == "__main__": unittest.main()