90 lines
2.8 KiB
Python
90 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
|
|
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 # 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 ModelScopeDiscoveryTests(unittest.TestCase):
|
|
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_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()
|