gabriel / musehub public
test_wire_streaming_push.py python
223 lines 7.8 KB
Raw
sha256:a34090cc4a394a78bd72cbbe34b08cc59525141e19135b6c0ab154f10611b9ef debug(push/stream): instrument O-frame decode path with INF… Sonnet 4.6 patch 121 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, now_utc_iso
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
63
64
65 def _header_frame(n_objects: int = 0, n_commits: int = 1) -> bytes:
66 return _wrap(SFRAME_HEADER, {
67 "t": SFRAME_HEADER,
68 "branch": "main",
69 "force": False,
70 "have": [],
71 "head": blob_id(b"head"),
72 "n_objects": n_objects,
73 "n_commits": n_commits,
74 })
75
76
77 def _commit_pack_frame(commits: list[dict], snapshots: list[dict]) -> bytes:
78 return _wrap(SFRAME_COMMIT_PACK, {"t": SFRAME_COMMIT_PACK, "commits": commits, "snapshots": snapshots})
79
80
81 def _end_frame(n_objects: int = 0, n_commits: int = 1) -> bytes:
82 return _wrap(SFRAME_END, {"t": SFRAME_END, "n_objects": n_objects, "n_commits": n_commits})
83
84
85 def _make_commit(snapshot_id: str | None = None, branch: str = "main") -> JSONObject:
86 snap = snapshot_id or blob_id(b"default-snap")
87 return {
88 "commit_id": blob_id(f"commit-{now_utc_iso()}".encode()),
89 "parent_ids": [],
90 "snapshot_id": snap,
91 "branch": branch,
92 "message": "test commit",
93 "author": "gabriel",
94 "committed_at": now_utc_iso(),
95 "signature": "",
96 "signer_key_id": "",
97 "agent_id": "",
98 "model_id": "",
99 "metadata": {},
100 }
101
102
103 def _make_snapshot(snap_id: str) -> JSONObject:
104 return {"snapshot_id": snap_id, "manifest": {}}
105
106
107 def _stub_r2_backend(monkeypatch: pytest.MonkeyPatch) -> None:
108 backend = MagicMock()
109 backend.store_object = AsyncMock(return_value=None)
110 backend.object_exists = AsyncMock(return_value=False)
111 backend.get_object = AsyncMock(return_value=None)
112 monkeypatch.setattr("musehub.services.musehub_wire.get_backend", lambda: backend)
113
114
115 # ---------------------------------------------------------------------------
116 # S1 — route source must NOT contain `await request.body()`
117 # ---------------------------------------------------------------------------
118
119 def test_s1_push_stream_route_does_not_call_request_body() -> None:
120 """push_stream must use request.stream(), not await request.body().
121
122 `await request.body()` forces Cloudflare to buffer the entire upload
123 before forwarding it to the origin server — causing bad_record_mac
124 on large payloads. This is a static check on the route source code.
125 """
126 from musehub.api.routes import wire
127
128 source = inspect.getsource(wire.push_stream)
129 assert "await request.body()" not in source, (
130 "push_stream route calls `await request.body()` which buffers the full body "
131 "before processing. Replace with `request.stream()` so Cloudflare streams "
132 "chunks as they arrive. See docs/protocol/musewire-performance.md."
133 )
134
135
136 # ---------------------------------------------------------------------------
137 # S2 — route source must call `request.stream()`
138 # ---------------------------------------------------------------------------
139
140 def test_s2_push_stream_route_uses_request_stream() -> None:
141 """push_stream must read the request body as a stream, not buffer it.
142
143 Accepts either request.stream() or the ASGI body_iter pattern — both stream
144 chunks as they arrive rather than buffering the full upload.
145 """
146 from musehub.api.routes import wire
147
148 source = inspect.getsource(wire.push_stream)
149 streams = "request.stream()" in source or "body_iter" in source
150 assert streams, (
151 "push_stream route must stream the request body (not buffer via request.body()). "
152 "Use request.stream() or an ASGI body_iter that reads from receive()."
153 )
154
155
156 # ---------------------------------------------------------------------------
157 # S3 — valid push still succeeds after streaming fix
158 # ---------------------------------------------------------------------------
159
160 @pytest.mark.asyncio
161 async def test_s3_valid_push_still_succeeds_with_streaming(
162 client: AsyncClient,
163 db_session: AsyncSession,
164 auth_headers: StrDict,
165 monkeypatch: pytest.MonkeyPatch,
166 ) -> None:
167 """After switching to request.stream(), a valid push must still yield RESULT ok=True."""
168 _stub_r2_backend(monkeypatch)
169 repo = await create_repo(db_session, owner="testuser", name="s3-success")
170
171 snap_id = blob_id(b"s3-snap")
172 commit = _make_commit(snapshot_id=snap_id)
173 snap = _make_snapshot(snap_id)
174 body = _header_frame() + _commit_pack_frame([commit], [snap]) + _end_frame()
175
176 resp = await client.post(
177 f"/{repo.owner}/{repo.slug}/push/stream",
178 content=body,
179 headers={**auth_headers, "Content-Type": WIRE_CONTENT_TYPE},
180 )
181
182 assert resp.status_code == 200
183 result = _last_frame(resp.content)
184 assert result.get("ok") is True, f"Expected ok=True, got: {result}"
185
186
187 # ---------------------------------------------------------------------------
188 # S4 — auth is checked before stream consumption
189 # ---------------------------------------------------------------------------
190
191 @pytest.mark.asyncio
192 async def test_s4_auth_checked_before_stream_consumed(
193 client: AsyncClient,
194 db_session: AsyncSession,
195 monkeypatch: pytest.MonkeyPatch,
196 ) -> None:
197 """Unauthenticated push must be rejected without consuming the request stream."""
198 _stub_r2_backend(monkeypatch)
199 repo = await create_repo(db_session, owner="testuser", name="s4-auth")
200
201 wire_was_called = []
202
203 async def _spy(session: AsyncSession, repo_id: str, body_iter: AsyncIterator[bytes], pusher_id: str | None) -> None:
204 wire_was_called.append(True)
205 yield _pack({"t": SFRAME_RESULT, "ok": True, "msg": "ok", "heads": {}, "head": ""})
206
207 monkeypatch.setattr("musehub.api.routes.wire.wire_push_stream", _spy)
208
209 body = _header_frame() + _end_frame()
210 resp = await client.post(
211 f"/{repo.owner}/{repo.slug}/push/stream",
212 content=body,
213 headers={"Content-Type": WIRE_CONTENT_TYPE},
214 )
215
216 if resp.status_code in (401, 403):
217 assert not wire_was_called, "wire_push_stream called despite auth failure"
218 else:
219 unpacker = msgpack.Unpacker(raw=False)
220 unpacker.feed(resp.content)
221 frames = list(unpacker)
222 error_frames = [f for f in frames if f.get("t") == "X"]
223 assert error_frames or not wire_was_called
File History 1 commit
sha256:a34090cc4a394a78bd72cbbe34b08cc59525141e19135b6c0ab154f10611b9ef debug(push/stream): instrument O-frame decode path with INF… Sonnet 4.6 patch 121 days ago