submmit/tests/test_modelscope_discovery.py

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  # 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_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_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()