116 lines
4.4 KiB
Python
116 lines
4.4 KiB
Python
"""Tests for Hugging Face auth on the llama.cpp prebuilt installer's fetches.
|
|
|
|
Anonymous huggingface.co downloads (tiny GGUF validation model) share a
|
|
per-IP rate limit that CI fleets exhaust (HTTP 429), forcing the prebuilt
|
|
path into a source build. auth_headers now sends HF_TOKEN to huggingface.co
|
|
hosts, and a redirect handler strips Authorization when the download is
|
|
redirected to a different host (CDN signed URLs). All tests are offline.
|
|
"""
|
|
|
|
import importlib.util
|
|
import sys
|
|
import urllib.request
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
|
|
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
|
|
|
|
_MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
|
|
_SPEC = importlib.util.spec_from_file_location(
|
|
"studio_install_llama_prebuilt_hf_auth", _MODULE_PATH
|
|
)
|
|
assert _SPEC is not None and _SPEC.loader is not None
|
|
mod = importlib.util.module_from_spec(_SPEC)
|
|
sys.modules[_SPEC.name] = mod
|
|
_SPEC.loader.exec_module(mod)
|
|
|
|
_TOKEN_VARS = ("GH_TOKEN", "GITHUB_TOKEN", "HF_TOKEN", "HUGGING_FACE_HUB_TOKEN")
|
|
HF_URL = (
|
|
"https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf"
|
|
)
|
|
GH_URL = "https://api.github.com/repos/unslothai/llama.cpp/releases"
|
|
|
|
|
|
def _headers(url, env):
|
|
"""auth_headers under a fully controlled token environment."""
|
|
with patch.dict(mod.os.environ, env, clear = False):
|
|
for var in _TOKEN_VARS:
|
|
if var not in env:
|
|
mod.os.environ.pop(var, None)
|
|
return mod.auth_headers(url)
|
|
|
|
|
|
class TestAuthHeaderRouting:
|
|
def test_hf_token_sent_to_huggingface(self):
|
|
headers = _headers(HF_URL, {"HF_TOKEN": "hf_x"})
|
|
assert headers.get("Authorization") == "Bearer hf_x"
|
|
|
|
def test_hub_token_fallback(self):
|
|
headers = _headers(HF_URL, {"HUGGING_FACE_HUB_TOKEN": "hf_y"})
|
|
assert headers.get("Authorization") == "Bearer hf_y"
|
|
|
|
def test_hf_token_not_sent_to_github(self):
|
|
headers = _headers(GH_URL, {"HF_TOKEN": "hf_x"})
|
|
assert "Authorization" not in headers
|
|
|
|
def test_hf_token_not_sent_to_other_hosts(self):
|
|
headers = _headers("https://cdn-lfs.huggingface.co/x", {"HF_TOKEN": "hf_x"})
|
|
assert "Authorization" not in headers
|
|
|
|
def test_gh_token_not_sent_to_huggingface(self):
|
|
headers = _headers(HF_URL, {"GH_TOKEN": "gh_x"})
|
|
assert "Authorization" not in headers
|
|
|
|
def test_gh_token_still_wins_on_github(self):
|
|
headers = _headers(GH_URL, {"GH_TOKEN": "gh_x", "HF_TOKEN": "hf_x"})
|
|
assert headers.get("Authorization") == "Bearer gh_x"
|
|
|
|
def test_no_tokens_no_auth(self):
|
|
assert "Authorization" not in _headers(HF_URL, {})
|
|
|
|
def test_validation_model_url_is_hf(self):
|
|
assert mod.should_send_hf_auth(mod.TEST_MODEL_URL) is True
|
|
|
|
|
|
class TestCrossHostRedirectStripsAuth:
|
|
def _redirect(self, newurl):
|
|
req = urllib.request.Request(HF_URL, headers = {"Authorization": "Bearer hf_x"})
|
|
handler = mod._CrossHostAuthStrippingRedirectHandler()
|
|
return handler.redirect_request(req, None, 302, "Found", {}, newurl)
|
|
|
|
def test_cross_host_redirect_drops_authorization(self):
|
|
new_request = self._redirect("https://cdn-lfs.huggingface.co/signed/blob")
|
|
assert new_request is not None
|
|
assert "Authorization" not in new_request.headers
|
|
assert "Authorization" not in new_request.unredirected_hdrs
|
|
|
|
def test_same_host_redirect_keeps_authorization(self):
|
|
new_request = self._redirect("https://huggingface.co/elsewhere/blob")
|
|
assert new_request is not None
|
|
assert new_request.headers.get("Authorization") == "Bearer hf_x"
|
|
|
|
|
|
class TestDownloadBytesWiring:
|
|
def test_download_bytes_sends_hf_auth(self):
|
|
response = MagicMock()
|
|
response.__enter__ = lambda s: s
|
|
response.__exit__ = lambda s, *a: False
|
|
response.headers.get.return_value = None
|
|
response.read.side_effect = [b"data", b""]
|
|
with (
|
|
patch.object(mod._URL_OPENER, "open", return_value = response) as opened,
|
|
patch.dict(mod.os.environ, {"HF_TOKEN": "hf_x"}, clear = False),
|
|
):
|
|
for var in ("GH_TOKEN", "GITHUB_TOKEN"):
|
|
mod.os.environ.pop(var, None)
|
|
data = mod.download_bytes(HF_URL)
|
|
assert data == b"data"
|
|
request = opened.call_args.args[0]
|
|
assert request.headers.get("Authorization") == "Bearer hf_x"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main([__file__, "-q"]))
|