141 lines
4.7 KiB
Python
141 lines
4.7 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
from datetime import UTC, datetime
|
|
from types import SimpleNamespace
|
|
|
|
import httpx
|
|
import pytest
|
|
from huggingface_hub.errors import (
|
|
GatedRepoError,
|
|
RepositoryNotFoundError,
|
|
RevisionNotFoundError,
|
|
)
|
|
|
|
from modelforge_api.providers.huggingface import (
|
|
HuggingFaceGated,
|
|
HuggingFaceNotFound,
|
|
HuggingFaceRevisionNotFound,
|
|
HuggingFaceUnavailable,
|
|
OfficialHuggingFaceProvider,
|
|
classify_file,
|
|
)
|
|
|
|
|
|
def response(status: int = 404) -> httpx.Response:
|
|
request = httpx.Request("GET", "https://huggingface.co/api/models/org/model")
|
|
return httpx.Response(status, request=request)
|
|
|
|
|
|
def test_search_uses_official_client_and_preserves_unknowns() -> None:
|
|
provider = OfficialHuggingFaceProvider()
|
|
calls: dict = {}
|
|
|
|
def list_models(**kwargs):
|
|
calls.update(kwargs)
|
|
return [
|
|
SimpleNamespace(
|
|
id="org/model",
|
|
sha="a" * 40,
|
|
gated=False,
|
|
private=False,
|
|
pipeline_tag=None,
|
|
library_name=None,
|
|
tags=None,
|
|
downloads=3,
|
|
likes=1,
|
|
last_modified=datetime.now(UTC),
|
|
)
|
|
]
|
|
|
|
provider.api.list_models = list_models # type: ignore[method-assign]
|
|
result = provider.search("model", limit=5, sort="likes", pipeline_tag=None)
|
|
assert result[0]["repository_id"] == "org/model"
|
|
assert result[0]["pipeline_tag"] is None
|
|
assert calls["search"] == "model" and calls["full"] is True
|
|
|
|
|
|
def test_snapshot_resolves_sha_inventory_card_and_scanner_evidence() -> None:
|
|
provider = OfficialHuggingFaceProvider()
|
|
provider.api.model_info = lambda *args, **kwargs: SimpleNamespace( # type: ignore[method-assign]
|
|
id="org/model",
|
|
sha="b" * 40,
|
|
gated=False,
|
|
private=False,
|
|
author="org",
|
|
pipeline_tag="feature-extraction",
|
|
library_name="transformers",
|
|
tags=["safetensors"],
|
|
downloads=10,
|
|
likes=2,
|
|
card_data={"license": "apache-2.0"},
|
|
security_repo_status={"status": "safe"},
|
|
last_modified=datetime.now(UTC),
|
|
siblings=[
|
|
SimpleNamespace(
|
|
rfilename="model.safetensors",
|
|
size=12,
|
|
blob_id="blob",
|
|
lfs={"sha256": "c" * 64, "size": 12, "pointer_size": 128},
|
|
)
|
|
],
|
|
)
|
|
snapshot = provider.snapshot("org/model", "main")
|
|
assert snapshot.resolved_commit_sha == "b" * 40
|
|
assert snapshot.files[0].upstream_sha256 == "c" * 64
|
|
assert snapshot.card_metadata["license"] == "apache-2.0"
|
|
assert snapshot.security_metadata["evidence_only"] is True
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("error", "expected"),
|
|
[
|
|
(GatedRepoError("gated", response=response(403)), HuggingFaceGated),
|
|
(RevisionNotFoundError("revision", response=response()), HuggingFaceRevisionNotFound),
|
|
(RepositoryNotFoundError("missing", response=response()), HuggingFaceNotFound),
|
|
(httpx.ReadTimeout("timeout", request=response().request), HuggingFaceUnavailable),
|
|
],
|
|
)
|
|
def test_snapshot_normalizes_provider_failures(error: Exception, expected: type[Exception]) -> None:
|
|
provider = OfficialHuggingFaceProvider()
|
|
|
|
def fail(*_args, **_kwargs):
|
|
raise error
|
|
|
|
provider.api.model_info = fail # type: ignore[method-assign]
|
|
with pytest.raises(expected):
|
|
provider.snapshot("org/model", "main")
|
|
|
|
|
|
def test_missing_exact_sha_fails_closed() -> None:
|
|
provider = OfficialHuggingFaceProvider()
|
|
provider.api.model_info = lambda *args, **kwargs: SimpleNamespace( # type: ignore[method-assign]
|
|
id="org/model",
|
|
sha=None,
|
|
gated=False,
|
|
private=False,
|
|
siblings=[],
|
|
card_data=None,
|
|
security_repo_status=None,
|
|
last_modified=None,
|
|
)
|
|
with pytest.raises(HuggingFaceUnavailable, match="exact commit"):
|
|
provider.snapshot("org/model", "main")
|
|
|
|
|
|
def test_file_classification_exposes_pickle_remote_code_and_unknown() -> None:
|
|
assert classify_file("model.safetensors") == ("safetensors", "weights", ())
|
|
assert "pickle_or_executable_serialization" in classify_file("weights.bin")[2]
|
|
assert classify_file("modeling_custom.py")[2] == ("remote_code",)
|
|
assert classify_file("LICENSE")[0] == "unknown"
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
os.getenv("MODELFORGE_RUN_LIVE_HF_TESTS") != "1",
|
|
reason="opt-in live Hub test; offline release gate uses deterministic provider tests",
|
|
)
|
|
def test_live_public_metadata_smoke() -> None:
|
|
snapshot = OfficialHuggingFaceProvider(timeout=30).snapshot("Qwen/Qwen3-Embedding-0.6B", "main")
|
|
assert len(snapshot.resolved_commit_sha) >= 40
|
|
assert snapshot.files
|