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