a6dc10dad2* Reduce and tighten comments and docstrings in tests Shorten verbose comments and docstrings across the test suite without changing any test logic. Remove narration that restates the next line, collapse long module and test docstrings to a single line, and drop banner separators. Keep regression context (issue and PR references, run ids), skip reasons, mocking and timing rationale, license headers, lint and type directives, and commented-out code. Comments and docstrings only: an AST signature check confirms no code, assertions, or string literals changed, and the suite byte-compiles cleanly. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
107 lines
4.1 KiB
Python
107 lines
4.1 KiB
Python
"""HF auth on the llama.cpp prebuilt installer: auth_headers sends HF_TOKEN to huggingface.co only, and a redirect handler strips Authorization on cross-host redirects. 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):
|
|
"""Call 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"]))
|