documentlanguagemodel Public

Watch 0 Fork 0 Star 0
Python · 13210 bytes Raw Blame History
1 """Unit coverage for the share pull orchestrator."""
2
3 from __future__ import annotations
4
5 import importlib
6 import sys
7 from pathlib import Path
8 from types import ModuleType, SimpleNamespace
9 from typing import cast
10
11 import pytest
12
13 from dlm.share.errors import ShareError, SinkError
14 from dlm.share.pull import (
15 PullResult,
16 _dispatch_pull,
17 _log_verification,
18 _try_hf_sidecar,
19 _try_peer_sidecar,
20 _try_url_sidecar,
21 pull,
22 )
23 from dlm.share.signing import VerifyResult, VerifyStatus
24 from dlm.share.sinks import SinkKind, SinkSpec
25
26 pull_mod = importlib.import_module("dlm.share.pull")
27
28
29 class TestPull:
30 def test_pull_dispatches_verifies_and_unpacks(
31 self,
32 tmp_path: Path,
33 monkeypatch: pytest.MonkeyPatch,
34 ) -> None:
35 source = "https://example.test/adapter.dlm.pack"
36 out_dir = tmp_path / "out"
37 home = tmp_path / "home"
38 progress = object()
39 spec = SinkSpec(kind=SinkKind.URL, target=source)
40 order: list[str] = []
41 verification = VerifyResult(status=VerifyStatus.VERIFIED, key_path=tmp_path / "trusted.pub")
42
43 monkeypatch.setattr(
44 pull_mod, "parse_source", lambda value: spec if value == source else None
45 )
46
47 def _fake_dispatch(
48 actual_spec: SinkSpec,
49 pack_path: Path,
50 sig_path: Path,
51 *,
52 progress: object | None,
53 ) -> int:
54 order.append("dispatch")
55 assert actual_spec == spec
56 assert pack_path.name == "incoming.dlm.pack"
57 assert sig_path.name == "incoming.dlm.pack.minisig"
58 assert progress is not None
59 pack_path.write_bytes(b"pack-bytes")
60 sig_path.write_text("signature", encoding="utf-8")
61 return 123
62
63 def _fake_verify(pack_path: Path, sig_path: Path) -> VerifyResult:
64 order.append("verify")
65 assert pack_path.read_bytes() == b"pack-bytes"
66 assert sig_path.read_text(encoding="utf-8") == "signature"
67 return verification
68
69 def _fake_unpack(
70 pack_path: Path,
71 *,
72 home: Path | None,
73 force: bool,
74 out_dir: Path,
75 ) -> SimpleNamespace:
76 order.append("unpack")
77 assert pack_path.read_bytes() == b"pack-bytes"
78 assert home == tmp_path / "home"
79 assert force is True
80 assert out_dir == tmp_path / "out"
81 return SimpleNamespace(
82 dlm_path=out_dir / "restored.dlm",
83 store_path=home / "store" / "01HZPULL",
84 dlm_id="01HZPULL",
85 )
86
87 monkeypatch.setattr(pull_mod, "_dispatch_pull", _fake_dispatch)
88 monkeypatch.setattr(pull_mod, "verify_signature", _fake_verify)
89 monkeypatch.setattr(pull_mod, "pack_unpack", _fake_unpack)
90
91 result = pull(
92 source,
93 out_dir=out_dir,
94 force=True,
95 home=home,
96 progress=cast("object", progress),
97 )
98
99 assert result == PullResult(
100 dlm_path=out_dir / "restored.dlm",
101 store_path=home / "store" / "01HZPULL",
102 dlm_id="01HZPULL",
103 source=source,
104 bytes_received=123,
105 verification=verification,
106 )
107 assert order == ["dispatch", "verify", "unpack"]
108
109
110 class TestDispatchPull:
111 def test_dispatch_pull_hf_downloads_pack_and_sidecar(
112 self,
113 tmp_path: Path,
114 monkeypatch: pytest.MonkeyPatch,
115 ) -> None:
116 import dlm.share.hf_sink as hf_sink
117
118 pack_path = tmp_path / "pack.dlm.pack"
119 sig_path = tmp_path / "pack.dlm.pack.minisig"
120 progress = object()
121 seen: dict[str, object] = {}
122
123 def _fake_pull_hf(repo_id: str, out_path: Path, *, progress: object | None = None) -> int:
124 seen["repo_id"] = repo_id
125 seen["progress"] = progress
126 out_path.write_bytes(b"hf-pack")
127 return 7
128
129 monkeypatch.setattr(hf_sink, "pull_hf", _fake_pull_hf)
130 monkeypatch.setattr(
131 pull_mod,
132 "_try_hf_sidecar",
133 lambda repo_id, sidecar_path: seen.update(
134 {"sidecar_repo_id": repo_id, "sidecar_path": sidecar_path}
135 ),
136 )
137
138 bytes_received = _dispatch_pull(
139 SinkSpec(kind=SinkKind.HF, target="org/repo"),
140 pack_path,
141 sig_path,
142 progress=cast("object", progress),
143 )
144
145 assert bytes_received == 7
146 assert pack_path.read_bytes() == b"hf-pack"
147 assert seen == {
148 "repo_id": "org/repo",
149 "progress": progress,
150 "sidecar_repo_id": "org/repo",
151 "sidecar_path": sig_path,
152 }
153
154 def test_dispatch_pull_url_downloads_pack_and_sidecar(
155 self,
156 tmp_path: Path,
157 monkeypatch: pytest.MonkeyPatch,
158 ) -> None:
159 import dlm.share.url_sink as url_sink
160
161 pack_path = tmp_path / "pack.dlm.pack"
162 sig_path = tmp_path / "pack.dlm.pack.minisig"
163 seen: dict[str, object] = {}
164
165 def _fake_pull_url(url: str, out_path: Path, *, progress: object | None = None) -> int:
166 seen["url"] = url
167 seen["progress"] = progress
168 out_path.write_bytes(b"url-pack")
169 return 9
170
171 monkeypatch.setattr(url_sink, "pull_url", _fake_pull_url)
172 monkeypatch.setattr(
173 pull_mod,
174 "_try_url_sidecar",
175 lambda url, sidecar_path: seen.update(
176 {"sidecar_url": url, "sidecar_path": sidecar_path}
177 ),
178 )
179
180 bytes_received = _dispatch_pull(
181 SinkSpec(kind=SinkKind.URL, target="https://example.test/pack"),
182 pack_path,
183 sig_path,
184 progress=None,
185 )
186
187 assert bytes_received == 9
188 assert pack_path.read_bytes() == b"url-pack"
189 assert seen == {
190 "url": "https://example.test/pack",
191 "progress": None,
192 "sidecar_url": "https://example.test/pack",
193 "sidecar_path": sig_path,
194 }
195
196 def test_dispatch_pull_peer_downloads_pack_and_sidecar(
197 self,
198 tmp_path: Path,
199 monkeypatch: pytest.MonkeyPatch,
200 ) -> None:
201 import dlm.share.peer as peer
202
203 pack_path = tmp_path / "pack.dlm.pack"
204 sig_path = tmp_path / "pack.dlm.pack.minisig"
205 seen: dict[str, object] = {}
206
207 def _fake_pull_peer(target: str, out_path: Path, *, progress: object | None = None) -> int:
208 seen["target"] = target
209 seen["progress"] = progress
210 out_path.write_bytes(b"peer-pack")
211 return 11
212
213 monkeypatch.setattr(peer, "pull_peer", _fake_pull_peer)
214 monkeypatch.setattr(
215 pull_mod,
216 "_try_peer_sidecar",
217 lambda target, sidecar_path: seen.update(
218 {"sidecar_target": target, "sidecar_path": sidecar_path}
219 ),
220 )
221
222 bytes_received = _dispatch_pull(
223 SinkSpec(kind=SinkKind.PEER, target="host:7337/pack?token=abc"),
224 pack_path,
225 sig_path,
226 progress=None,
227 )
228
229 assert bytes_received == 11
230 assert pack_path.read_bytes() == b"peer-pack"
231 assert seen == {
232 "target": "host:7337/pack?token=abc",
233 "progress": None,
234 "sidecar_target": "host:7337/pack?token=abc",
235 "sidecar_path": sig_path,
236 }
237
238 def test_dispatch_pull_local_copies_pack_and_signature(self, tmp_path: Path) -> None:
239 src = tmp_path / "src.dlm.pack"
240 sig = tmp_path / "src.dlm.pack.minisig"
241 src.write_bytes(b"local-pack")
242 sig.write_text("local-signature", encoding="utf-8")
243 pack_path = tmp_path / "incoming.dlm.pack"
244 sig_path = tmp_path / "incoming.dlm.pack.minisig"
245
246 bytes_received = _dispatch_pull(
247 SinkSpec(kind=SinkKind.LOCAL, target=str(src)),
248 pack_path,
249 sig_path,
250 progress=None,
251 )
252
253 assert bytes_received == len(b"local-pack")
254 assert pack_path.read_bytes() == b"local-pack"
255 assert sig_path.read_text(encoding="utf-8") == "local-signature"
256
257 def test_dispatch_pull_local_missing_source_raises(self, tmp_path: Path) -> None:
258 with pytest.raises(SinkError, match="source missing"):
259 _dispatch_pull(
260 SinkSpec(kind=SinkKind.LOCAL, target=str(tmp_path / "missing.dlm.pack")),
261 tmp_path / "incoming.dlm.pack",
262 tmp_path / "incoming.dlm.pack.minisig",
263 progress=None,
264 )
265
266 def test_dispatch_pull_rejects_unsupported_kind(self, tmp_path: Path) -> None:
267 weird = SinkSpec(kind=cast("SinkKind", "weird"), target="x")
268
269 with pytest.raises(ShareError, match="unsupported sink kind"):
270 _dispatch_pull(
271 weird,
272 tmp_path / "incoming.dlm.pack",
273 tmp_path / "incoming.dlm.pack.minisig",
274 progress=None,
275 )
276
277
278 class TestPullSidecars:
279 def test_try_hf_sidecar_copies_downloaded_signature(
280 self,
281 tmp_path: Path,
282 monkeypatch: pytest.MonkeyPatch,
283 ) -> None:
284 fake_hub = ModuleType("huggingface_hub")
285 fake_utils = ModuleType("huggingface_hub.utils")
286 downloaded = tmp_path / "downloaded.minisig"
287 downloaded.write_text("hf-signature", encoding="utf-8")
288
289 class FakeHfHubHTTPError(Exception):
290 pass
291
292 def _fake_download(*, repo_id: str, filename: str, repo_type: str) -> str:
293 assert repo_id == "org/repo"
294 assert filename == "adapter.dlm.pack.minisig"
295 assert repo_type == "model"
296 return str(downloaded)
297
298 fake_hub.hf_hub_download = _fake_download
299 fake_utils.HfHubHTTPError = FakeHfHubHTTPError
300 monkeypatch.setitem(sys.modules, "huggingface_hub", fake_hub)
301 monkeypatch.setitem(sys.modules, "huggingface_hub.utils", fake_utils)
302
303 sig_path = tmp_path / "incoming.minisig"
304 _try_hf_sidecar("org/repo", sig_path)
305
306 assert sig_path.read_text(encoding="utf-8") == "hf-signature"
307
308 def test_try_hf_sidecar_suppresses_hub_errors(
309 self,
310 tmp_path: Path,
311 monkeypatch: pytest.MonkeyPatch,
312 ) -> None:
313 fake_hub = ModuleType("huggingface_hub")
314 fake_utils = ModuleType("huggingface_hub.utils")
315
316 class FakeHfHubHTTPError(Exception):
317 pass
318
319 def _fake_download(*, repo_id: str, filename: str, repo_type: str) -> str:
320 raise FakeHfHubHTTPError("missing")
321
322 fake_hub.hf_hub_download = _fake_download
323 fake_utils.HfHubHTTPError = FakeHfHubHTTPError
324 monkeypatch.setitem(sys.modules, "huggingface_hub", fake_hub)
325 monkeypatch.setitem(sys.modules, "huggingface_hub.utils", fake_utils)
326
327 sig_path = tmp_path / "incoming.minisig"
328 _try_hf_sidecar("org/repo", sig_path)
329
330 assert not sig_path.exists()
331
332 def test_try_url_sidecar_suppresses_missing_sidecar(
333 self,
334 tmp_path: Path,
335 monkeypatch: pytest.MonkeyPatch,
336 ) -> None:
337 import dlm.share.url_sink as url_sink
338
339 def _fake_pull_url(url: str, out_path: Path, *, progress: object | None = None) -> int:
340 raise SinkError(f"missing {url}")
341
342 monkeypatch.setattr(url_sink, "pull_url", _fake_pull_url)
343
344 _try_url_sidecar("https://example.test/pack", tmp_path / "incoming.minisig")
345
346 def test_try_peer_sidecar_suppresses_missing_sidecar(
347 self,
348 tmp_path: Path,
349 monkeypatch: pytest.MonkeyPatch,
350 ) -> None:
351 import dlm.share.peer as peer
352
353 def _fake_pull_peer(target: str, out_path: Path, *, progress: object | None = None) -> int:
354 raise SinkError(f"missing {target}")
355
356 monkeypatch.setattr(peer, "pull_peer", _fake_pull_peer)
357
358 _try_peer_sidecar("host:7337/pack?token=abc", tmp_path / "incoming.minisig")
359
360
361 class TestVerificationLogging:
362 def test_log_verification_verified(
363 self, caplog: pytest.LogCaptureFixture, tmp_path: Path
364 ) -> None:
365 caplog.set_level("INFO")
366
367 _log_verification(
368 "hf:org/repo",
369 VerifyResult(status=VerifyStatus.VERIFIED, key_path=tmp_path / "trusted.pub"),
370 )
371
372 assert "verified signature" in caplog.text
373
374 def test_log_verification_unverified(self, caplog: pytest.LogCaptureFixture) -> None:
375 caplog.set_level("WARNING")
376
377 _log_verification(
378 "https://example.test/pack",
379 VerifyResult(status=VerifyStatus.UNVERIFIED, detail="no trusted key matched"),
380 )
381
382 assert "signature present but could not verify" in caplog.text
383 assert "no trusted key matched" in caplog.text
384
385 def test_log_verification_unsigned(self, caplog: pytest.LogCaptureFixture) -> None:
386 caplog.set_level("INFO")
387
388 _log_verification(
389 "./local.dlm.pack",
390 VerifyResult(status=VerifyStatus.UNSIGNED),
391 )
392
393 assert "no signature" in caplog.text