"""TDD — fetch client presign routing: stream vs presigned R2. The presigned fetch path mirrors the presigned push path: - Small fetches (< 500 objects AND < 50 MB) go through fetch/stream as before. - Large fetches call POST /fetch/presign first; if the server returns presign=True it downloads the bundle directly from the presigned URL, bypassing Cloudflare entirely. - LocalFileTransport never presigns (loopback / local dev). Test plan --------- Unit / integration FPR0 Below object threshold → _use_fetch_presign returns False. FPR1 At object threshold, non-loopback → _use_fetch_presign returns True. FPR2 At byte threshold (via size hint from server), non-loopback → True. FPR3 Loopback URL (file://) → always False regardless of counts. FPR4 Server returns presign=False (LocalBackend on staging) → falls through to fetch/stream normally; no crash. FPR5 Server returns presign=True → client downloads bundle from presigned URL, no call to fetch/stream endpoint. FPR6 Presigned download writes correct objects (bundle bytes parsed and dispatched through on_object callback). FPR7 fetch_presign_or_stream is the new single entry-point: small repo → delegates to fetch_stream unchanged. FPR8 fetch_presign_or_stream: large repo, server presigns → returns same FetchStreamResult shape as fetch_stream. FPR9 Network error on presigned GET → raises TransportError (not silent drop). FPR10 _FETCH_PRESIGN_OBJECT_THRESHOLD and _FETCH_PRESIGN_BYTE_THRESHOLD match the server-side constants. Security FPRS0 LocalFileTransport.fetch_presign_or_stream never touches the presign endpoint — always delegates to fetch_stream unconditionally. FPRS1 file:// and localhost URLs are classified as loopback by routing logic. Stress / state integrity FPRST0 One of N parallel GETs returns non-200 → TransportError raised. FPRST1 N=10 parallel GETs all succeed → all objects dispatched via on_object. Performance FPRP0 on_object receives every object when presign=True (no silent drops). FPRP1 on_object is never called when presign_response has empty object_urls. """ from __future__ import annotations from unittest.mock import AsyncMock, MagicMock, patch, call import io import msgpack import pytest import pathlib from muse.core.types import fake_id from muse.core.paths import muse_dir from muse.core.pack import ObjectPayload type ObjectMap = dict[str, bytes] type ObjectUrlMap = dict[str, str] # --------------------------------------------------------------------------- # FPR10 — constants match server side (verified first so mismatches are obvious) # --------------------------------------------------------------------------- def test_fpr10_threshold_constants_match_server() -> None: """Client-side thresholds must equal server FETCH_PRESIGN_*_THRESHOLD.""" from muse.cli.commands.pull import ( _FETCH_PRESIGN_OBJECT_THRESHOLD, _FETCH_PRESIGN_BYTE_THRESHOLD, ) assert _FETCH_PRESIGN_OBJECT_THRESHOLD == 500 assert _FETCH_PRESIGN_BYTE_THRESHOLD == 50 * 1024 * 1024 # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- _SERVER_OBJECT_THRESHOLD = 500 _SERVER_BYTE_THRESHOLD = 50 * 1024 * 1024 def _is_loopback(url: str) -> bool: return url.startswith("file://") or "localhost" in url or "127.0.0.1" in url def _use_fetch_presign(n_objects: int, total_bytes: int, is_loopback: bool) -> bool: """Mirror the client routing decision from pull.py.""" from muse.cli.commands.pull import ( _FETCH_PRESIGN_OBJECT_THRESHOLD, _FETCH_PRESIGN_BYTE_THRESHOLD, ) return ( not is_loopback and ( n_objects >= _FETCH_PRESIGN_OBJECT_THRESHOLD or total_bytes >= _FETCH_PRESIGN_BYTE_THRESHOLD ) ) _REPO_ID = fake_id("repo") _COMMIT_ID = fake_id("commit") _SNAP_ID = fake_id("snap") def _make_objects(n: int) -> ObjectMap: """Return {oid: raw_bytes} for n objects.""" return {fake_id(f"obj-{i}"): f"object-{i}".encode() for i in range(n)} def _make_presign_response(object_urls: ObjectUrlMap, n_objects: int = 0) -> bytes: """Build a presign=True msgpack response with per-object URLs.""" manifest = {f"file_{i}.py": oid for i, oid in enumerate(object_urls)} return msgpack.packb({ "presign": True, "object_urls": object_urls, "commits": [{ "commit_id": _COMMIT_ID, "snapshot_id": _SNAP_ID, "message": "test", "author": "gabriel", "committed_at": "2026-04-30T00:00:00+00:00", "parent_commit_id": None, "parent2_commit_id": None, "agent_id": "", "model_id": "", "toolchain_id": "", "signer_public_key": "", "signature": "", }], "snapshots": [{"snapshot_id": _SNAP_ID, "manifest": manifest, "directories": [], "created_at": ""}], "branch_heads": {"main": _COMMIT_ID}, "repo_id": _REPO_ID, "domain": "code", "default_branch": "main", "expires_at": "2026-04-30T02:00:00+00:00", "commit_count": 1, "object_count": n_objects or len(object_urls), }, use_bin_type=True) # --------------------------------------------------------------------------- # FPR0 — below threshold → stream path # --------------------------------------------------------------------------- def test_fpr0_below_threshold_uses_stream() -> None: assert not _use_fetch_presign( n_objects=_SERVER_OBJECT_THRESHOLD - 1, total_bytes=0, is_loopback=False, ) # --------------------------------------------------------------------------- # FPR1 — at object threshold, non-loopback → presign # --------------------------------------------------------------------------- def test_fpr1_object_threshold_triggers_presign() -> None: assert _use_fetch_presign( n_objects=_SERVER_OBJECT_THRESHOLD, total_bytes=0, is_loopback=False, ) # --------------------------------------------------------------------------- # FPR2 — at byte threshold → presign # --------------------------------------------------------------------------- def test_fpr2_byte_threshold_triggers_presign() -> None: assert _use_fetch_presign( n_objects=1, total_bytes=_SERVER_BYTE_THRESHOLD, is_loopback=False, ) # --------------------------------------------------------------------------- # FPR3 — loopback URL never presigns # --------------------------------------------------------------------------- def test_fpr3_loopback_never_presigns() -> None: assert not _use_fetch_presign( n_objects=_SERVER_OBJECT_THRESHOLD * 10, total_bytes=_SERVER_BYTE_THRESHOLD * 10, is_loopback=True, ) # --------------------------------------------------------------------------- # FPR4 — server returns presign=False → fallthrough to fetch/stream # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_fpr4_server_presign_false_falls_through_to_stream() -> None: """When server says presign=False, client must call fetch/stream normally.""" from muse.core.transport import HttpTransport, FetchStreamResult transport = HttpTransport() url = "https://staging.musehub.ai/gabriel/muse" want = [fake_id("want")] have: list[str] = [] presign_response = msgpack.packb({ "presign": False, "object_count": 50, "commit_count": 10, }, use_bin_type=True) fake_stream_result: FetchStreamResult = FetchStreamResult( repo_id=fake_id("repo"), domain="code", default_branch="main", branch_heads={"main": want[0]}, commits=[], snapshots=[], objects_received=0, shallow_commits=[], ) calls: list[str] = [] class _FakeResp: def __init__(self, body: bytes, status: int = 200) -> None: self.status_code = status self.content = body with patch.object(transport, "fetch_stream", return_value=fake_stream_result) as mock_stream, \ patch("httpx.Client") as mock_client_cls: mock_client = MagicMock() mock_client.__enter__ = MagicMock(return_value=mock_client) mock_client.__exit__ = MagicMock(return_value=False) mock_client.post = MagicMock(return_value=_FakeResp(presign_response)) mock_client_cls.return_value = mock_client result = transport.fetch_presign_or_stream( url, None, want=want, have=have, on_object=None, ) mock_stream.assert_called_once_with(url, None, want=want, have=have, on_object=None) assert result is fake_stream_result # --------------------------------------------------------------------------- # FPR5 — server returns presign=True → download from URL, no fetch/stream # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_fpr5_server_presign_true_skips_stream() -> None: """When server returns presign=True, client downloads per-object URLs; fetch/stream not called.""" from muse.core.transport import HttpTransport transport = HttpTransport() url = "https://staging.musehub.ai/gabriel/muse" want = [fake_id("want")] have: list[str] = [] objects = _make_objects(2) object_urls = {oid: f"https://r2.example.com/{oid}?sig=x" for oid in objects} presign_response = _make_presign_response(object_urls) class _FakeResp: def __init__(self, body: bytes, status: int = 200) -> None: self.status_code = status self.content = body stream_called = [] with patch.object(transport, "fetch_stream", side_effect=lambda *a, **kw: stream_called.append(1)), \ patch("httpx.Client") as mock_client_cls: mock_client = MagicMock() mock_client.__enter__ = MagicMock(return_value=mock_client) mock_client.__exit__ = MagicMock(return_value=False) mock_client.post = MagicMock(return_value=_FakeResp(presign_response)) mock_client.get = MagicMock(side_effect=lambda u, **kw: _FakeResp( objects.get(next((oid for oid in objects if oid in u), ""), b"raw-content") )) mock_client_cls.return_value = mock_client result = transport.fetch_presign_or_stream( url, None, want=want, have=have, on_object=None, ) assert not stream_called, "fetch/stream must NOT be called when presign=True" assert result["objects_received"] == 2 assert result["commit_count"] == 1 # --------------------------------------------------------------------------- # FPR6 — presigned bundle bytes dispatched via on_object callback # --------------------------------------------------------------------------- def test_fpr6_presigned_objects_dispatched_via_on_object() -> None: """Raw bytes from each per-object GET are dispatched through on_object.""" from muse.core.transport import HttpTransport from muse.core.pack import ObjectPayload transport = HttpTransport() url = "https://staging.musehub.ai/gabriel/muse" want = [fake_id("want")] objects = _make_objects(3) object_urls = {oid: f"https://r2.example.com/{oid}?sig=x" for oid in objects} presign_response = _make_presign_response(object_urls) received: list[ObjectPayload] = [] def _on_object(obj: ObjectPayload) -> None: received.append(obj) class _FakeResp: def __init__(self, body: bytes, status: int = 200) -> None: self.status_code = status self.content = body def _fake_get(u: str) -> _FakeResp: for oid, content in objects.items(): if oid in u: return _FakeResp(content) return _FakeResp(b"unexpected") with patch("httpx.Client") as mock_client_cls, \ patch.object(transport, "fetch_stream"): mock_client = MagicMock() mock_client.__enter__ = MagicMock(return_value=mock_client) mock_client.__exit__ = MagicMock(return_value=False) mock_client.post = MagicMock(return_value=_FakeResp(presign_response)) mock_client.get = MagicMock(side_effect=_fake_get) mock_client_cls.return_value = mock_client transport.fetch_presign_or_stream( url, None, want=want, have=[], on_object=_on_object, ) assert len(received) == 3 for obj in received: assert obj["object_id"].startswith("sha256:") assert obj["content"] # --------------------------------------------------------------------------- # FPR7 — small repo → fetch_presign_or_stream delegates to fetch_stream # --------------------------------------------------------------------------- def test_fpr7_small_repo_delegates_to_fetch_stream() -> None: """fetch_presign_or_stream with small remote info → pure fetch_stream delegation.""" from muse.core.transport import HttpTransport, FetchStreamResult transport = HttpTransport() url = "https://staging.musehub.ai/gabriel/timing-test" want = [fake_id("want")] have: list[str] = [] # Server says presign=False (small repo, below threshold) presign_response = msgpack.packb({ "presign": False, "object_count": 42, "commit_count": 10, }, use_bin_type=True) fake_result: FetchStreamResult = FetchStreamResult( repo_id=fake_id("repo"), domain="code", default_branch="main", branch_heads={"main": want[0]}, commits=[], snapshots=[], objects_received=42, shallow_commits=[], ) class _FakeResp: def __init__(self, body: bytes, status: int = 200) -> None: self.status_code = status self.content = body with patch.object(transport, "fetch_stream", return_value=fake_result) as mock_stream, \ patch("httpx.Client") as mock_client_cls: mock_client = MagicMock() mock_client.__enter__ = MagicMock(return_value=mock_client) mock_client.__exit__ = MagicMock(return_value=False) mock_client.post = MagicMock(return_value=_FakeResp(presign_response)) mock_client_cls.return_value = mock_client result = transport.fetch_presign_or_stream(url, None, want=want, have=have) mock_stream.assert_called_once() assert result["objects_received"] == 42 # --------------------------------------------------------------------------- # FPR8 — large repo, presign=True → FetchStreamResult shape correct # --------------------------------------------------------------------------- def test_fpr8_large_repo_presign_returns_correct_shape() -> None: """fetch_presign_or_stream returns FetchStreamResult-compatible dict when presign=True.""" from muse.core.transport import HttpTransport transport = HttpTransport() url = "https://staging.musehub.ai/gabriel/muse" want = [fake_id("want")] objects = _make_objects(1) object_urls = {oid: f"https://r2.example.com/{oid}?sig=abc" for oid in objects} presign_response = _make_presign_response(object_urls) class _FakeResp: def __init__(self, body: bytes, status: int = 200) -> None: self.status_code = status self.content = body def _fake_get(u: str) -> _FakeResp: for oid, content in objects.items(): if oid in u: return _FakeResp(content) return _FakeResp(b"unexpected") with patch("httpx.Client") as mock_client_cls, \ patch.object(transport, "fetch_stream"): mock_client = MagicMock() mock_client.__enter__ = MagicMock(return_value=mock_client) mock_client.__exit__ = MagicMock(return_value=False) mock_client.post = MagicMock(return_value=_FakeResp(presign_response)) mock_client.get = MagicMock(side_effect=_fake_get) mock_client_cls.return_value = mock_client result = transport.fetch_presign_or_stream(url, None, want=want, have=[]) # Must have same top-level keys as FetchStreamResult for key in ("repo_id", "domain", "default_branch", "branch_heads", "commits", "snapshots", "objects_received", "shallow_commits"): assert key in result, f"missing key: {key}" assert result["objects_received"] == 1 assert result["commit_count"] == 1 # --------------------------------------------------------------------------- # FPR9 — network error on presigned GET → TransportError # --------------------------------------------------------------------------- def test_fpr9_presigned_get_error_raises_transport_error() -> None: """A non-200 response from the presigned URL raises TransportError.""" from muse.core.transport import HttpTransport, TransportError transport = HttpTransport() url = "https://staging.musehub.ai/gabriel/muse" want = [fake_id("want")] objects = _make_objects(1) object_urls = {oid: f"https://r2.example.com/{oid}?sig=abc" for oid in objects} presign_response = _make_presign_response(object_urls) class _FakeResp: def __init__(self, body: bytes, status: int = 200) -> None: self.status_code = status self.content = body with patch("httpx.Client") as mock_client_cls, \ patch.object(transport, "fetch_stream"): mock_client = MagicMock() mock_client.__enter__ = MagicMock(return_value=mock_client) mock_client.__exit__ = MagicMock(return_value=False) mock_client.post = MagicMock(return_value=_FakeResp(presign_response)) mock_client.get = MagicMock(return_value=_FakeResp(b"Access Denied", status=403)) mock_client_cls.return_value = mock_client with pytest.raises(TransportError, match="403"): transport.fetch_presign_or_stream(url, None, want=want, have=[]) # =========================================================================== # Security tests # =========================================================================== # --------------------------------------------------------------------------- # FPRS0 — LocalFileTransport never calls the presign endpoint # --------------------------------------------------------------------------- def test_fprs0_local_transport_never_hits_presign_endpoint(tmp_path: pathlib.Path) -> None: """LocalFileTransport.fetch_presign_or_stream must delegate to fetch_stream, not presign.""" from muse.core.transport import LocalFileTransport, FetchStreamResult import pathlib dot_muse = muse_dir(tmp_path) dot_muse.mkdir() url = f"file://{tmp_path}" fake_result: FetchStreamResult = FetchStreamResult( repo_id=fake_id("repo"), domain="code", default_branch="main", branch_heads={}, commits=[], snapshots=[], objects_received=0, shallow_commits=[], ) transport = LocalFileTransport() with patch.object(transport, "fetch_stream", return_value=fake_result) as mock_stream, \ patch("httpx.Client") as mock_http: result = transport.fetch_presign_or_stream( url, None, want=[fake_id("w")], have=[], ) mock_stream.assert_called_once() mock_http.assert_not_called() assert result is fake_result # --------------------------------------------------------------------------- # FPRS1 — file:// and localhost URLs classified as loopback # --------------------------------------------------------------------------- def test_fprs1_loopback_url_classification() -> None: """file://, localhost, and 127.0.0.1 must be classified as loopback.""" assert _is_loopback("file:///home/gabriel/repo") assert _is_loopback("https://localhost:1337/gabriel/muse") assert _is_loopback("http://127.0.0.1:8000/gabriel/muse") assert not _is_loopback("https://staging.musehub.ai/gabriel/muse") assert not _is_loopback("https://musehub.ai/gabriel/muse") # =========================================================================== # Stress / state integrity tests # =========================================================================== # --------------------------------------------------------------------------- # FPRST0 — one of N parallel GETs fails → TransportError # --------------------------------------------------------------------------- def test_fprst0_one_failing_get_raises_transport_error() -> None: """If any parallel presigned GET returns non-200, TransportError is raised.""" from muse.core.transport import HttpTransport, TransportError transport = HttpTransport() url = "https://staging.musehub.ai/gabriel/muse" want = [fake_id("want")] n = 5 objects = _make_objects(n) object_urls = {oid: f"https://r2.example.com/{oid}?sig=st0" for oid in objects} presign_response = _make_presign_response(object_urls) call_count = 0 class _FakeResp: def __init__(self, body: bytes, status: int = 200) -> None: self.status_code = status self.content = body def _failing_get(u: str) -> _FakeResp: nonlocal call_count call_count += 1 # Fail on the third request if call_count == 3: return _FakeResp(b"Internal Server Error", status=500) for oid, content in objects.items(): if oid in u: return _FakeResp(content) return _FakeResp(b"not found", status=404) with patch("httpx.Client") as mock_client_cls, \ patch.object(transport, "fetch_stream"): mock_client = MagicMock() mock_client.__enter__ = MagicMock(return_value=mock_client) mock_client.__exit__ = MagicMock(return_value=False) mock_client.post = MagicMock(return_value=_FakeResp(presign_response)) mock_client.get = MagicMock(side_effect=_failing_get) mock_client_cls.return_value = mock_client with pytest.raises(TransportError): transport.fetch_presign_or_stream(url, None, want=want, have=[]) # --------------------------------------------------------------------------- # FPRST1 — N=10 parallel GETs succeed → all dispatched via on_object # --------------------------------------------------------------------------- def test_fprst1_n_parallel_gets_all_dispatched() -> None: """When all N presigned GETs succeed, on_object receives exactly N payloads.""" from muse.core.transport import HttpTransport from muse.core.pack import ObjectPayload n = 10 transport = HttpTransport() url = "https://staging.musehub.ai/gabriel/muse" want = [fake_id("want")] objects = _make_objects(n) object_urls = {oid: f"https://r2.example.com/{oid}?sig=st1" for oid in objects} presign_response = _make_presign_response(object_urls) received: list[ObjectPayload] = [] class _FakeResp: def __init__(self, body: bytes, status: int = 200) -> None: self.status_code = status self.content = body def _fake_get(u: str) -> _FakeResp: for oid, content in objects.items(): if oid in u: return _FakeResp(content) return _FakeResp(b"unexpected") with patch("httpx.Client") as mock_client_cls, \ patch.object(transport, "fetch_stream"): mock_client = MagicMock() mock_client.__enter__ = MagicMock(return_value=mock_client) mock_client.__exit__ = MagicMock(return_value=False) mock_client.post = MagicMock(return_value=_FakeResp(presign_response)) mock_client.get = MagicMock(side_effect=_fake_get) mock_client_cls.return_value = mock_client result = transport.fetch_presign_or_stream( url, None, want=want, have=[], on_object=received.append, ) assert len(received) == n assert result["objects_received"] == n # =========================================================================== # Performance tests # =========================================================================== # --------------------------------------------------------------------------- # FPRP0 — on_object receives every object when presign=True # --------------------------------------------------------------------------- def test_fprp0_on_object_receives_all_presigned_objects() -> None: """on_object callback must be invoked once per object in the presigned map.""" from muse.core.transport import HttpTransport from muse.core.pack import ObjectPayload transport = HttpTransport() url = "https://staging.musehub.ai/gabriel/muse" want = [fake_id("want")] objects = _make_objects(8) object_urls = {oid: f"https://r2.example.com/{oid}?sig=p0" for oid in objects} presign_response = _make_presign_response(object_urls) dispatched_ids: set[str] = set() class _FakeResp: def __init__(self, body: bytes, status: int = 200) -> None: self.status_code = status self.content = body def _fake_get(u: str) -> _FakeResp: for oid, content in objects.items(): if oid in u: return _FakeResp(content) return _FakeResp(b"unexpected") def _on_object(obj: ObjectPayload) -> None: dispatched_ids.add(obj["object_id"]) with patch("httpx.Client") as mock_client_cls, \ patch.object(transport, "fetch_stream"): mock_client = MagicMock() mock_client.__enter__ = MagicMock(return_value=mock_client) mock_client.__exit__ = MagicMock(return_value=False) mock_client.post = MagicMock(return_value=_FakeResp(presign_response)) mock_client.get = MagicMock(side_effect=_fake_get) mock_client_cls.return_value = mock_client transport.fetch_presign_or_stream(url, None, want=want, have=[], on_object=_on_object) assert dispatched_ids == set(objects.keys()) # --------------------------------------------------------------------------- # FPRP1 — on_object never called when object_urls is empty # --------------------------------------------------------------------------- def test_fprp1_no_on_object_calls_for_empty_presign_response() -> None: """When server returns presign=False (object_count=0), on_object must not be called.""" from muse.core.transport import HttpTransport, FetchStreamResult transport = HttpTransport() url = "https://staging.musehub.ai/gabriel/muse" want = [fake_id("want")] presign_response = msgpack.packb({ "presign": False, "object_count": 0, "commit_count": 0, }, use_bin_type=True) fake_result: FetchStreamResult = FetchStreamResult( repo_id=fake_id("repo"), domain="code", default_branch="main", branch_heads={}, commits=[], snapshots=[], objects_received=0, shallow_commits=[], ) on_object_calls: list[ObjectPayload] = [] class _FakeResp: def __init__(self, body: bytes, status: int = 200) -> None: self.status_code = status self.content = body with patch.object(transport, "fetch_stream", return_value=fake_result), \ patch("httpx.Client") as mock_client_cls: mock_client = MagicMock() mock_client.__enter__ = MagicMock(return_value=mock_client) mock_client.__exit__ = MagicMock(return_value=False) mock_client.post = MagicMock(return_value=_FakeResp(presign_response)) mock_client_cls.return_value = mock_client transport.fetch_presign_or_stream( url, None, want=want, have=[], on_object=on_object_calls.append, ) assert on_object_calls == [], "on_object must not be called when presign=False"