gabriel / musehub public
test_wire_oc_server.py python
337 lines 11.4 KB
Raw
sha256:7b5741c91286e310b18bdae2435d6740836ac7fb4d33f31857e0181a16d6faf9 feat(wire): Phase 14B server-side OC-frame reassembly Sonnet 4.6 patch 154 days ago
1 """TDD — Phase 14B server-side: OC-frame reassembly.
2
3 Rules encoded here:
4
5 P14B-5 Server assembles a 1-chunk OC sequence → stored object matches original.
6
7 P14B-6 Server assembles a 3-chunk OC sequence → stored object matches original.
8
9 P14B-7 Server receives partial OC sequence (2 of 3 chunks) then END frame →
10 returns an error frame; does not store a partial object.
11 """
12 from __future__ import annotations
13
14 import hashlib
15 from unittest.mock import AsyncMock, MagicMock
16
17 import msgpack
18 import pytest
19 from sqlalchemy.ext.asyncio import AsyncSession
20
21 from muse.core.mpack import MuseWireFrameWriter, grpc_frame
22 from musehub.models.wire import (
23 SFRAME_HEADER, SFRAME_OBJECT_CHUNK, SFRAME_COMMIT_PACK, SFRAME_END,
24 SFRAME_ERROR, SFRAME_RESULT,
25 )
26
27
28 # ---------------------------------------------------------------------------
29 # Shared helpers (same pattern as test_wire_push_stream.py)
30 # ---------------------------------------------------------------------------
31
32 _fw = MuseWireFrameWriter()
33
34
35 def _pack(data: object) -> bytes:
36 return msgpack.packb(data, use_bin_type=True)
37
38
39 def _wrap(ft: str, data: object) -> bytes:
40 return grpc_frame(_fw.wrap(frame_type=ft, payload=_pack(data)))
41
42
43 def _sha256_oid(raw: bytes) -> str:
44 return "sha256:" + hashlib.sha256(raw).hexdigest()
45
46
47 def _header_frame(n_objects: int = 1) -> bytes:
48 return _wrap(SFRAME_HEADER, {
49 "t": SFRAME_HEADER, "branch": "main", "force": False, "have": [],
50 "head": "sha256:" + "a" * 64, "n_objects": n_objects, "n_commits": 0,
51 "agent_id": "", "model_id": "", "sig": "",
52 })
53
54
55 def _oc_frame(oid: str, ci: int, tc: int, chunk: bytes, *,
56 enc: str | None = None, path: str | None = None,
57 sz: int | None = None) -> bytes:
58 payload: dict = {"t": SFRAME_OBJECT_CHUNK, "id": oid,
59 "ci": ci, "tc": tc, "content": chunk}
60 if enc is not None:
61 payload["enc"] = enc
62 if path is not None:
63 payload["path"] = path
64 if sz is not None:
65 payload["sz"] = sz
66 return _wrap(SFRAME_OBJECT_CHUNK, payload)
67
68
69 def _c_frame() -> bytes:
70 return _wrap(SFRAME_COMMIT_PACK, {
71 "t": SFRAME_COMMIT_PACK, "commits": [], "snapshots": [],
72 "snapshot_deltas": [],
73 })
74
75
76 def _e_frame() -> bytes:
77 return _wrap(SFRAME_END, {"t": SFRAME_END})
78
79
80 async def _collect_frames(gen) -> list[dict]:
81 import struct
82 results = []
83 async for frame_bytes in gen:
84 if len(frame_bytes) < 5:
85 continue
86 msg_len = struct.unpack(">I", frame_bytes[1:5])[0]
87 raw = frame_bytes[5: 5 + msg_len]
88 try:
89 results.append(msgpack.unpackb(raw, raw=False))
90 except Exception:
91 pass
92 return results
93
94
95 # ---------------------------------------------------------------------------
96 # Shared fixtures
97 # ---------------------------------------------------------------------------
98
99 @pytest.fixture()
100 def stored_objects():
101 return {}
102
103
104 @pytest.fixture()
105 def stub_backend(monkeypatch, stored_objects):
106 backend = AsyncMock()
107 backend.exists = AsyncMock(side_effect=lambda oid: oid in stored_objects)
108 def _put(oid, data, path=""):
109 stored_objects[oid] = data
110 return f"local://{oid}"
111 backend.put = AsyncMock(side_effect=_put)
112 backend.get = AsyncMock(side_effect=lambda oid: stored_objects.get(oid))
113 monkeypatch.setattr("musehub.services.musehub_wire.get_backend", lambda: backend)
114 return backend
115
116
117 @pytest.fixture()
118 def stub_session():
119 session = AsyncMock(spec=AsyncSession)
120 session.execute = AsyncMock(
121 return_value=MagicMock(scalar=lambda: None, fetchall=lambda: [])
122 )
123 session.commit = AsyncMock()
124 session.add = MagicMock()
125 return session
126
127
128 # ---------------------------------------------------------------------------
129 # P14B-5 — Server assembles 1-chunk OC sequence → stored correctly
130 # ---------------------------------------------------------------------------
131
132 class TestP14B5OneChunkOC:
133 @pytest.mark.asyncio
134 async def test_single_oc_chunk_stored_correctly(
135 self, stub_backend, stub_session, stored_objects
136 ) -> None:
137 from musehub.services.musehub_wire import wire_push_stream
138
139 content = b"hello world from a single OC chunk"
140 oid = _sha256_oid(content)
141
142 async def body():
143 yield (
144 _header_frame(n_objects=1)
145 + _oc_frame(oid, 0, 1, content, enc="raw", path="file.txt", sz=len(content))
146 + _c_frame()
147 + _e_frame()
148 )
149
150 await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
151 assert oid in stored_objects, f"Object {oid[:20]} not stored"
152 assert stored_objects[oid] == content
153
154 @pytest.mark.asyncio
155 async def test_single_oc_no_error_frame(
156 self, stub_backend, stub_session
157 ) -> None:
158 from musehub.services.musehub_wire import wire_push_stream
159
160 content = b"no error expected"
161 oid = _sha256_oid(content)
162
163 async def body():
164 yield (
165 _header_frame(n_objects=1)
166 + _oc_frame(oid, 0, 1, content, enc="raw", path="", sz=len(content))
167 + _c_frame()
168 + _e_frame()
169 )
170
171 frames = await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
172 assert not [f for f in frames if f.get("t") == "X"]
173
174 @pytest.mark.asyncio
175 async def test_result_ok_after_single_oc(
176 self, stub_backend, stub_session
177 ) -> None:
178 from musehub.services.musehub_wire import wire_push_stream
179
180 content = b"result test"
181 oid = _sha256_oid(content)
182
183 async def body():
184 yield (
185 _header_frame(n_objects=1)
186 + _oc_frame(oid, 0, 1, content, enc="raw", path="", sz=len(content))
187 + _c_frame()
188 + _e_frame()
189 )
190
191 frames = await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
192 result_frames = [f for f in frames if f.get("t") == "R"]
193 assert result_frames, "No RESULT frame"
194 assert result_frames[0]["ok"] is True
195
196
197 # ---------------------------------------------------------------------------
198 # P14B-6 — Server assembles 3-chunk OC sequence → stored correctly
199 # ---------------------------------------------------------------------------
200
201 class TestP14B6ThreeChunkOC:
202 def _content_and_chunks(self, n: int = 900):
203 content = bytes([i % 256 for i in range(n)])
204 oid = _sha256_oid(content)
205 size = n // 3
206 chunks = [content[i * size:(i + 1) * size if i < 2 else n] for i in range(3)]
207 return content, oid, chunks
208
209 @pytest.mark.asyncio
210 async def test_three_chunks_stored_correctly(
211 self, stub_backend, stub_session, stored_objects
212 ) -> None:
213 from musehub.services.musehub_wire import wire_push_stream
214
215 content, oid, chunks = self._content_and_chunks()
216
217 async def body():
218 b = _header_frame(n_objects=1)
219 for i, chunk in enumerate(chunks):
220 kw = {"enc": "raw", "path": "big.bin", "sz": len(content)} if i == 0 else {}
221 b += _oc_frame(oid, i, 3, chunk, **kw)
222 b += _c_frame() + _e_frame()
223 yield b
224
225 await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
226 assert oid in stored_objects
227 assert stored_objects[oid] == content
228
229 @pytest.mark.asyncio
230 async def test_three_chunks_no_error_frame(
231 self, stub_backend, stub_session
232 ) -> None:
233 from musehub.services.musehub_wire import wire_push_stream
234
235 content, oid, chunks = self._content_and_chunks()
236
237 async def body():
238 b = _header_frame(n_objects=1)
239 for i, chunk in enumerate(chunks):
240 kw = {"enc": "raw", "path": "", "sz": len(content)} if i == 0 else {}
241 b += _oc_frame(oid, i, 3, chunk, **kw)
242 b += _c_frame() + _e_frame()
243 yield b
244
245 frames = await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
246 assert not [f for f in frames if f.get("t") == "X"]
247
248 @pytest.mark.asyncio
249 async def test_chunks_out_of_order_assembles_correctly(
250 self, stub_backend, stub_session, stored_objects
251 ) -> None:
252 from musehub.services.musehub_wire import wire_push_stream
253
254 content, oid, chunks = self._content_and_chunks()
255
256 async def body():
257 b = _header_frame(n_objects=1)
258 # Send in reverse order
259 for i in reversed(range(3)):
260 kw = {"enc": "raw", "path": "", "sz": len(content)} if i == 0 else {}
261 b += _oc_frame(oid, i, 3, chunks[i], **kw)
262 b += _c_frame() + _e_frame()
263 yield b
264
265 await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
266 assert oid in stored_objects
267 assert stored_objects[oid] == content
268
269
270 # ---------------------------------------------------------------------------
271 # P14B-7 — Partial OC sequence + END → error frame, no partial object stored
272 # ---------------------------------------------------------------------------
273
274 class TestP14B7PartialOCSequenceErrors:
275 @pytest.mark.asyncio
276 async def test_one_of_three_chunks_returns_error(
277 self, stub_backend, stub_session, stored_objects
278 ) -> None:
279 from musehub.services.musehub_wire import wire_push_stream
280
281 content = bytes(300)
282 oid = _sha256_oid(content)
283
284 async def body():
285 yield (
286 _header_frame(n_objects=1)
287 + _oc_frame(oid, 0, 3, content[:100], enc="raw", path="", sz=300)
288 # chunks 1 and 2 missing
289 + _c_frame()
290 + _e_frame()
291 )
292
293 frames = await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
294 assert [f for f in frames if f.get("t") == "X"], "Expected error for incomplete OC"
295
296 @pytest.mark.asyncio
297 async def test_partial_oc_object_not_stored(
298 self, stub_backend, stub_session, stored_objects
299 ) -> None:
300 from musehub.services.musehub_wire import wire_push_stream
301
302 content = bytes(300)
303 oid = _sha256_oid(content)
304
305 async def body():
306 yield (
307 _header_frame(n_objects=1)
308 + _oc_frame(oid, 0, 3, content[:100], enc="raw", path="", sz=300)
309 + _c_frame()
310 + _e_frame()
311 )
312
313 await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
314 assert oid not in stored_objects, "Partial object must not be stored"
315
316 @pytest.mark.asyncio
317 async def test_two_of_three_chunks_returns_error(
318 self, stub_backend, stub_session, stored_objects
319 ) -> None:
320 from musehub.services.musehub_wire import wire_push_stream
321
322 content = bytes(300)
323 oid = _sha256_oid(content)
324
325 async def body():
326 yield (
327 _header_frame(n_objects=1)
328 + _oc_frame(oid, 0, 3, content[:100], enc="raw", path="", sz=300)
329 + _oc_frame(oid, 1, 3, content[100:200])
330 # chunk 2 missing
331 + _c_frame()
332 + _e_frame()
333 )
334
335 frames = await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
336 assert [f for f in frames if f.get("t") == "X"], "Expected error for missing last chunk"
337 assert oid not in stored_objects
File History 1 commit
sha256:7b5741c91286e310b18bdae2435d6740836ac7fb4d33f31857e0181a16d6faf9 feat(wire): Phase 14B server-side OC-frame reassembly Sonnet 4.6 patch 154 days ago