File size: 8,289 Bytes
201dc77 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 | """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) # dies mid-line
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)) # no dup, no gap, no split record
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" # resumed on a line boundary
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) # ignores Range
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 # served from cache
# moving ref is never cached
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 # second pass fully cached
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"))
|