| """hub.py against a mocked transport: no network, no real cache dirs.""" |
| from __future__ import annotations |
|
|
| import io |
| import json |
| import urllib.error |
|
|
| import pytest |
|
|
| from align.data import hub |
|
|
|
|
| class FakeResp(io.BytesIO): |
| def __init__(self, data: bytes, status: int = 200, headers: dict | None = None, fail_after: int | None = None): |
| super().__init__(data) |
| self.status = status |
| self.headers = headers or {} |
| self._fail_after = fail_after |
| self._served = 0 |
|
|
| def getcode(self): |
| return self.status |
|
|
| def read(self, n=-1): |
| if self._fail_after is not None and self._served >= self._fail_after: |
| raise ConnectionResetError("boom") |
| out = super().read(n if self._fail_after is None else min(n, self._fail_after - self._served)) |
| self._served += len(out) |
| return out |
|
|
|
|
| @pytest.fixture(autouse=True) |
| def _no_sleep(monkeypatch): |
| monkeypatch.setattr(hub, "_sleep", lambda s: None) |
| monkeypatch.setenv("HF_TOKEN", "tok123") |
|
|
|
|
| def test_url_builders(): |
| assert hub.resolve_url("a/b", "data/x y.jsonl", "abc") == ( |
| "https://huggingface.co/datasets/a/b/resolve/abc/data/x%20y.jsonl" |
| ) |
| assert hub.tree_url("a/b", "main").endswith("/api/datasets/a/b/tree/main?recursive=true") |
| u = hub.rows_url("a/b", "default", "train", 200) |
| assert u.startswith("https://datasets-server.huggingface.co/rows?") and "offset=200" in u and "length=100" in u |
|
|
|
|
| def test_token_file_and_env(tmp_path, monkeypatch): |
| monkeypatch.delenv("HF_TOKEN") |
| t = tmp_path / "token" |
| t.write_text(" filetok\n") |
| assert hub.read_token(t) == "filetok" |
| assert hub.read_token(tmp_path / "missing") is None |
| monkeypatch.setenv("HF_TOKEN", "envtok") |
| assert hub.read_token(t) == "envtok" |
|
|
|
|
| def test_token_only_sent_to_hub_hosts(monkeypatch): |
| seen = [] |
|
|
| def fake(req, timeout): |
| seen.append((req.full_url, req.headers.get("Authorization"))) |
| return FakeResp(b"x") |
|
|
| monkeypatch.setattr(hub, "_urlopen", fake) |
| hub._open_with_retry("https://huggingface.co/x") |
| hub._open_with_retry("https://cdn.example.com/x") |
| assert seen[0][1] == "Bearer tok123" and seen[1][1] is None |
|
|
|
|
| def test_redirect_strips_auth_across_hosts(): |
| import urllib.request |
|
|
| h = hub._StripAuthOnRedirect() |
| req = urllib.request.Request("https://huggingface.co/a", headers={"Authorization": "Bearer t"}) |
| new = h.redirect_request(req, None, 302, "Found", {}, "https://cdn.example.com/b") |
| assert "Authorization" not in {k.title() for k in new.headers} | {k.title() for k in new.unredirected_hdrs} |
| new2 = h.redirect_request(req, None, 302, "Found", {}, "https://huggingface.co/c") |
| assert "Authorization" in {k.title() for k in new2.headers} |
|
|
|
|
| def test_stream_resume_uses_range_at_line_boundary(monkeypatch): |
| payload = b"".join(json.dumps({"i": i}).encode() + b"\n" for i in range(20)) |
| calls = [] |
|
|
| def fake(req, timeout): |
| rng = req.headers.get("Range") |
| calls.append(rng) |
| if rng is None: |
| return FakeResp(payload, 200, fail_after=len(payload) // 2 + 3) |
| start = int(rng.split("=")[1].rstrip("-")) |
| return FakeResp(payload[start:], 206) |
|
|
| monkeypatch.setattr(hub, "_urlopen", fake) |
| got = [json.loads(x)["i"] for x in hub.iter_lines("https://huggingface.co/f")] |
| assert got == list(range(20)) |
| assert calls[0] is None and calls[1].startswith("bytes=") and len(calls) == 2 |
| start = int(calls[1].split("=")[1].rstrip("-")) |
| assert payload[start - 1 : start] == b"\n" |
|
|
|
|
| def test_resume_refuses_if_server_ignores_range(monkeypatch): |
| payload = b"a\nb\nc\nd\n" |
|
|
| def fake(req, timeout): |
| if req.headers.get("Range") is None: |
| return FakeResp(payload, 200, fail_after=5) |
| return FakeResp(payload, 200) |
|
|
| monkeypatch.setattr(hub, "_urlopen", fake) |
| with pytest.raises(hub.HubError, match="ignored Range"): |
| list(hub.iter_lines("https://huggingface.co/f")) |
|
|
|
|
| def test_iter_jsonl_skips_blank_and_no_trailing_newline(monkeypatch): |
| monkeypatch.setattr(hub, "_urlopen", lambda r, t: FakeResp(b'{"a":1}\n\n{"a":2}')) |
| assert list(hub.iter_jsonl("o/r", "f.jsonl")) == [{"a": 1}, {"a": 2}] |
|
|
|
|
| def test_retry_then_fail_fast_on_404(monkeypatch): |
| n = {"c": 0} |
|
|
| def flaky(req, timeout): |
| n["c"] += 1 |
| if n["c"] < 3: |
| raise urllib.error.HTTPError(req.full_url, 503, "busy", {}, None) |
| return FakeResp(b"ok") |
|
|
| monkeypatch.setattr(hub, "_urlopen", flaky) |
| assert hub._open_with_retry("https://huggingface.co/x").read() == b"ok" and n["c"] == 3 |
|
|
| def gone(req, timeout): |
| raise urllib.error.HTTPError(req.full_url, 404, "nope", {}, None) |
|
|
| monkeypatch.setattr(hub, "_urlopen", gone) |
| with pytest.raises(hub.HubError, match="404"): |
| hub._open_with_retry("https://huggingface.co/x") |
|
|
|
|
| def test_list_repo_files_follows_link_and_caches_immutable(monkeypatch, tmp_path): |
| sha = "a" * 40 |
| pages = { |
| hub.tree_url("o/r", sha): FakeResp( |
| json.dumps([{"path": "a.parquet", "type": "file"}]).encode(), |
| headers={"Link": '<https://huggingface.co/page2>; rel="next"'}, |
| ), |
| "https://huggingface.co/page2": FakeResp( |
| json.dumps([{"path": "d", "type": "directory"}, {"path": "b.jsonl", "type": "file"}]).encode() |
| ), |
| } |
| calls = [] |
|
|
| def fake(req, timeout): |
| calls.append(req.full_url) |
| return pages[req.full_url] |
|
|
| monkeypatch.setattr(hub, "_urlopen", fake) |
| assert hub.list_repo_file_paths("o/r", sha, cache_dir=tmp_path) == ["a.parquet", "b.jsonl"] |
| n = len(calls) |
| hub.list_repo_files("o/r", sha, cache_dir=tmp_path) |
| assert len(calls) == n |
| |
| monkeypatch.setattr(hub, "_urlopen", lambda r, t: FakeResp(b"[]")) |
| assert hub.list_repo_files("o/r", "main", cache_dir=tmp_path) == [] |
|
|
|
|
| def _page(offset, n, total, truncated=False): |
| rows = [ |
| {"row_idx": offset + i, "row": {"text": f"t{offset + i}"}, "truncated_cells": ["text"] if truncated else []} |
| for i in range(n) |
| ] |
| return json.dumps({"rows": rows, "num_rows_total": total}).encode() |
|
|
|
|
| def test_iter_rows_paginates_caches_and_respects_max(monkeypatch, tmp_path): |
| calls = [] |
|
|
| def fake(req, timeout): |
| calls.append(req.full_url) |
| off = int(req.full_url.split("offset=")[1].split("&")[0]) |
| return FakeResp(_page(off, min(100, 250 - off), 250)) |
|
|
| monkeypatch.setattr(hub, "_urlopen", fake) |
| rows = list(hub.iter_rows("o/r", "default", "train", cache_dir=tmp_path)) |
| assert [r["text"] for r in rows] == [f"t{i}" for i in range(250)] |
| assert len(calls) == 3 |
| list(hub.iter_rows("o/r", "default", "train", cache_dir=tmp_path)) |
| assert len(calls) == 3 |
| assert len(list(hub.iter_rows("o/r", max_rows=150, cache_dir=tmp_path / "c2"))) == 150 |
|
|
|
|
| def test_iter_rows_raises_on_truncated_cells(monkeypatch, tmp_path): |
| monkeypatch.setattr(hub, "_urlopen", lambda r, t: FakeResp(_page(0, 2, 2, truncated=True))) |
| with pytest.raises(hub.HubError, match="truncated"): |
| list(hub.iter_rows("o/r", cache_dir=tmp_path)) |
|
|
|
|
| def test_rows_length_bound(): |
| with pytest.raises(ValueError): |
| hub.fetch_rows_page("o/r", "default", "train", 0, length=101) |
|
|
|
|
| def test_corrupt_cache_page_is_refetched(monkeypatch, tmp_path): |
| url = hub.rows_url("o/r", "default", "train", 0, 100) |
| cf = hub._cache_file(tmp_path, "rows", url) |
| cf.parent.mkdir(parents=True) |
| cf.write_text("{not json") |
| monkeypatch.setattr(hub, "_urlopen", lambda r, t: FakeResp(_page(0, 1, 1))) |
| assert list(hub.iter_rows("o/r", cache_dir=tmp_path)) == [{"text": "t0"}] |
|
|
|
|
| def test_pretrain_stream_goes_through_hub(monkeypatch, tmp_path): |
| from align.data import pretrain_corpus as pc |
|
|
| monkeypatch.setattr(hub, "DEFAULT_CACHE_DIR", str(tmp_path)) |
| monkeypatch.setattr(hub, "_urlopen", lambda r, t: FakeResp(_page(0, 3, 3))) |
| assert list(pc.iter_streamed_texts("o/r")) == ["t0", "t1", "t2"] |
| with pytest.raises(KeyError, match="nope"): |
| list(pc.iter_streamed_texts("o/r", text_key="nope")) |
|
|
