gabriel / musehub public
test_wire_streaming_push.py python
228 lines 8.0 KB
Raw
sha256:bb2baaabdd19320bde50cb69d447fd1c1e571467df729be1a23b5e93064b046a feat(intel): standardize headers, gauge icon, velocity card… Sonnet 4.6 minor ⚠ breaking 142 days ago
1 """TDD — push_stream route must stream the request body (not buffer it).
2
3 Rules (from musewire-performance.md):
4
5 S1 push_stream route source must NOT contain `await request.body()`.
6 S2 push_stream route source must call `request.stream()`.
7 S3 A valid push still yields RESULT ok=True after the streaming change.
8 S4 Auth is checked before any stream bytes are consumed.
9
10 Why this matters: `await request.body()` causes Cloudflare to buffer the
11 entire request body before forwarding it to the origin, which triggers
12 TLS bad_record_mac errors on large payloads. request.stream() lets
13 Cloudflare forward chunks as they arrive — matching GitHub's push behavior.
14 """
15 from __future__ import annotations
16
17 import inspect
18 from collections.abc import AsyncIterator
19 from datetime import datetime, timezone
20 from unittest.mock import AsyncMock, MagicMock, patch
21
22 import msgpack
23 import pytest
24 from httpx import AsyncClient
25 from sqlalchemy.ext.asyncio import AsyncSession
26
27 from muse.core.types import blob_id
28 from musehub.types.json_types import JSONObject, JSONValue, StrDict
29 from musehub.models.wire import (
30 SFRAME_COMMIT_PACK,
31 SFRAME_END,
32 SFRAME_HEADER,
33 SFRAME_RESULT,
34 )
35 from muse.core.mpack import WIRE_CONTENT_TYPE, MuseWireFrameWriter
36 from tests.factories import create_repo
37
38 _fw = MuseWireFrameWriter()
39
40
41 def _last_frame(raw: bytes) -> JSONObject:
42 unpacker = msgpack.Unpacker(raw=False)
43 unpacker.feed(raw)
44 last: JSONObject = {}
45 for frame in unpacker:
46 last = frame
47 return last
48
49
50 # ---------------------------------------------------------------------------
51 # Helpers
52 # ---------------------------------------------------------------------------
53
54 def _pack(obj: JSONValue) -> bytes:
55 return msgpack.packb(obj, use_bin_type=True)
56
57
58 def _wrap(ft: str, data: JSONValue) -> bytes:
59 return _fw.wrap(frame_type=ft, payload=_pack(data))
60
61
62 def _sha256_oid(data: bytes) -> str:
63 return blob_id(data)
64
65
66 def _utc() -> str:
67 return datetime.now(tz=timezone.utc).isoformat()
68
69
70 def _header_frame(n_objects: int = 0, n_commits: int = 1) -> bytes:
71 return _wrap(SFRAME_HEADER, {
72 "t": SFRAME_HEADER,
73 "branch": "main",
74 "force": False,
75 "have": [],
76 "head": _sha256_oid(b"head"),
77 "n_objects": n_objects,
78 "n_commits": n_commits,
79 })
80
81
82 def _commit_pack_frame(commits: list[dict], snapshots: list[dict]) -> bytes:
83 return _wrap(SFRAME_COMMIT_PACK, {"t": SFRAME_COMMIT_PACK, "commits": commits, "snapshots": snapshots})
84
85
86 def _end_frame(n_objects: int = 0, n_commits: int = 1) -> bytes:
87 return _wrap(SFRAME_END, {"t": SFRAME_END, "n_objects": n_objects, "n_commits": n_commits})
88
89
90 def _make_commit(snapshot_id: str | None = None, branch: str = "main") -> JSONObject:
91 snap = snapshot_id or _sha256_oid(b"default-snap")
92 return {
93 "commit_id": _sha256_oid(f"commit-{_utc()}".encode()),
94 "parent_ids": [],
95 "snapshot_id": snap,
96 "branch": branch,
97 "message": "test commit",
98 "author": "gabriel",
99 "committed_at": _utc(),
100 "signature": "",
101 "signer_key_id": "",
102 "agent_id": "",
103 "model_id": "",
104 "metadata": {},
105 }
106
107
108 def _make_snapshot(snap_id: str) -> JSONObject:
109 return {"snapshot_id": snap_id, "manifest": {}}
110
111
112 def _stub_r2_backend(monkeypatch: pytest.MonkeyPatch) -> None:
113 backend = MagicMock()
114 backend.store_object = AsyncMock(return_value=None)
115 backend.object_exists = AsyncMock(return_value=False)
116 backend.get_object = AsyncMock(return_value=None)
117 monkeypatch.setattr("musehub.services.musehub_wire.get_backend", lambda: backend)
118
119
120 # ---------------------------------------------------------------------------
121 # S1 — route source must NOT contain `await request.body()`
122 # ---------------------------------------------------------------------------
123
124 def test_s1_push_stream_route_does_not_call_request_body() -> None:
125 """push_stream must use request.stream(), not await request.body().
126
127 `await request.body()` forces Cloudflare to buffer the entire upload
128 before forwarding it to the origin server — causing bad_record_mac
129 on large payloads. This is a static check on the route source code.
130 """
131 from musehub.api.routes import wire
132
133 source = inspect.getsource(wire.push_stream)
134 assert "await request.body()" not in source, (
135 "push_stream route calls `await request.body()` which buffers the full body "
136 "before processing. Replace with `request.stream()` so Cloudflare streams "
137 "chunks as they arrive. See docs/protocol/musewire-performance.md."
138 )
139
140
141 # ---------------------------------------------------------------------------
142 # S2 — route source must call `request.stream()`
143 # ---------------------------------------------------------------------------
144
145 def test_s2_push_stream_route_uses_request_stream() -> None:
146 """push_stream must read the request body as a stream, not buffer it.
147
148 Accepts either request.stream() or the ASGI body_iter pattern — both stream
149 chunks as they arrive rather than buffering the full upload.
150 """
151 from musehub.api.routes import wire
152
153 source = inspect.getsource(wire.push_stream)
154 streams = "request.stream()" in source or "body_iter" in source
155 assert streams, (
156 "push_stream route must stream the request body (not buffer via request.body()). "
157 "Use request.stream() or an ASGI body_iter that reads from receive()."
158 )
159
160
161 # ---------------------------------------------------------------------------
162 # S3 — valid push still succeeds after streaming fix
163 # ---------------------------------------------------------------------------
164
165 @pytest.mark.asyncio
166 async def test_s3_valid_push_still_succeeds_with_streaming(
167 client: AsyncClient,
168 db_session: AsyncSession,
169 auth_headers: StrDict,
170 monkeypatch: pytest.MonkeyPatch,
171 ) -> None:
172 """After switching to request.stream(), a valid push must still yield RESULT ok=True."""
173 _stub_r2_backend(monkeypatch)
174 repo = await create_repo(db_session, owner="testuser", name="s3-success")
175
176 snap_id = _sha256_oid(b"s3-snap")
177 commit = _make_commit(snapshot_id=snap_id)
178 snap = _make_snapshot(snap_id)
179 body = _header_frame() + _commit_pack_frame([commit], [snap]) + _end_frame()
180
181 resp = await client.post(
182 f"/{repo.owner}/{repo.slug}/push/stream",
183 content=body,
184 headers={**auth_headers, "Content-Type": WIRE_CONTENT_TYPE},
185 )
186
187 assert resp.status_code == 200
188 result = _last_frame(resp.content)
189 assert result.get("ok") is True, f"Expected ok=True, got: {result}"
190
191
192 # ---------------------------------------------------------------------------
193 # S4 — auth is checked before stream consumption
194 # ---------------------------------------------------------------------------
195
196 @pytest.mark.asyncio
197 async def test_s4_auth_checked_before_stream_consumed(
198 client: AsyncClient,
199 db_session: AsyncSession,
200 monkeypatch: pytest.MonkeyPatch,
201 ) -> None:
202 """Unauthenticated push must be rejected without consuming the request stream."""
203 _stub_r2_backend(monkeypatch)
204 repo = await create_repo(db_session, owner="testuser", name="s4-auth")
205
206 wire_was_called = []
207
208 async def _spy(session: AsyncSession, repo_id: str, body_iter: AsyncIterator[bytes], pusher_id: str | None) -> None:
209 wire_was_called.append(True)
210 yield _pack({"t": SFRAME_RESULT, "ok": True, "msg": "ok", "heads": {}, "head": ""})
211
212 monkeypatch.setattr("musehub.api.routes.wire.wire_push_stream", _spy)
213
214 body = _header_frame() + _end_frame()
215 resp = await client.post(
216 f"/{repo.owner}/{repo.slug}/push/stream",
217 content=body,
218 headers={"Content-Type": WIRE_CONTENT_TYPE},
219 )
220
221 if resp.status_code in (401, 403):
222 assert not wire_was_called, "wire_push_stream called despite auth failure"
223 else:
224 unpacker = msgpack.Unpacker(raw=False)
225 unpacker.feed(resp.content)
226 frames = list(unpacker)
227 error_frames = [f for f in frames if f.get("t") == "X"]
228 assert error_frames or not wire_was_called
File History 1 commit
sha256:bb2baaabdd19320bde50cb69d447fd1c1e571467df729be1a23b5e93064b046a feat(intel): standardize headers, gauge icon, velocity card… Sonnet 4.6 minor ⚠ 142 days ago