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