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