test_grpc_wire_framing.py
python
sha256:9590cee1e0ccd6c76528f005b95d634d80f5019f0dcb7c371e149adc31d1fb65
refactor: enforce gRPC framing on all MWP wire traffic
Sonnet 4.6
minor
⚠ breaking
155 days ago
| 1 | """Tests for gRPC framing strip on the server side. |
| 2 | |
| 3 | The server receives a stream of gRPC length-prefix frames. Each message is |
| 4 | one MWP frame. The server strips the 5-byte gRPC header, then feeds the |
| 5 | inner bytes to iter_wire_frames exactly as before. |
| 6 | |
| 7 | Test plan: |
| 8 | A. iter_wire_frames_grpc strips gRPC prefix and yields MWP frames correctly |
| 9 | B. Empty stream yields nothing |
| 10 | C. Multiple gRPC-wrapped MWP frames all decoded in order |
| 11 | D. gRPC frame with wrong compression flag raises WireFrameError |
| 12 | E. Truncated gRPC header raises WireFrameError |
| 13 | F. Truncated gRPC message body raises WireFrameError |
| 14 | G. Roundtrip: client grpc_frame(MuseWireFrameWriter.wrap(...)) → server iter |
| 15 | """ |
| 16 | from __future__ import annotations |
| 17 | |
| 18 | import struct |
| 19 | import asyncio |
| 20 | |
| 21 | import msgpack |
| 22 | import pytest |
| 23 | |
| 24 | |
| 25 | # --------------------------------------------------------------------------- |
| 26 | # Helpers |
| 27 | # --------------------------------------------------------------------------- |
| 28 | |
| 29 | def _pack(obj: object) -> bytes: |
| 30 | return msgpack.packb(obj, use_bin_type=True) |
| 31 | |
| 32 | |
| 33 | def _grpc_wrap(mwp_frame: bytes) -> bytes: |
| 34 | """Encode one MWP frame in gRPC length-prefix format.""" |
| 35 | return b"\x00" + struct.pack(">I", len(mwp_frame)) + mwp_frame |
| 36 | |
| 37 | |
| 38 | def _make_mwp_frame(ft: str, payload: bytes) -> bytes: |
| 39 | from muse.core.mpack import MuseWireFrameWriter |
| 40 | return MuseWireFrameWriter().wrap(frame_type=ft, payload=payload) |
| 41 | |
| 42 | |
| 43 | async def _collect(body: bytes) -> list[tuple[dict, bytes]]: |
| 44 | """Run iter_wire_frames_grpc over a byte blob, collect all (envelope, payload) pairs.""" |
| 45 | from musehub.services.musehub_wire import iter_wire_frames_grpc |
| 46 | |
| 47 | async def _iter(): |
| 48 | yield body |
| 49 | |
| 50 | results = [] |
| 51 | async for envelope, payload in iter_wire_frames_grpc(_iter()): |
| 52 | results.append((envelope, payload)) |
| 53 | return results |
| 54 | |
| 55 | |
| 56 | def _run(coro): |
| 57 | return asyncio.run(coro) |
| 58 | |
| 59 | |
| 60 | # --------------------------------------------------------------------------- |
| 61 | # A. Single frame roundtrip |
| 62 | # --------------------------------------------------------------------------- |
| 63 | |
| 64 | def test_single_frame_roundtrip(): |
| 65 | payload = _pack({"t": "E"}) |
| 66 | mwp = _make_mwp_frame("E", payload) |
| 67 | stream = _grpc_wrap(mwp) |
| 68 | results = _run(_collect(stream)) |
| 69 | assert len(results) == 1 |
| 70 | envelope, decoded_payload = results[0] |
| 71 | assert envelope["ft"] == "E" |
| 72 | assert decoded_payload == payload |
| 73 | |
| 74 | |
| 75 | # --------------------------------------------------------------------------- |
| 76 | # B. Empty stream |
| 77 | # --------------------------------------------------------------------------- |
| 78 | |
| 79 | def test_empty_stream_yields_nothing(): |
| 80 | results = _run(_collect(b"")) |
| 81 | assert results == [] |
| 82 | |
| 83 | |
| 84 | # --------------------------------------------------------------------------- |
| 85 | # C. Multiple frames in order |
| 86 | # --------------------------------------------------------------------------- |
| 87 | |
| 88 | def test_multiple_frames_decoded_in_order(): |
| 89 | frames = [ |
| 90 | _make_mwp_frame("H", _pack({"t": "H", "repo": "r", "branch": "b", |
| 91 | "have": [], "local_head": None, |
| 92 | "force": False, "object_count": 0})), |
| 93 | _make_mwp_frame("E", _pack({"t": "E"})), |
| 94 | ] |
| 95 | stream = b"".join(_grpc_wrap(f) for f in frames) |
| 96 | results = _run(_collect(stream)) |
| 97 | assert len(results) == 2 |
| 98 | assert results[0][0]["ft"] == "H" |
| 99 | assert results[1][0]["ft"] == "E" |
| 100 | |
| 101 | |
| 102 | # --------------------------------------------------------------------------- |
| 103 | # D. Wrong compression flag raises WireFrameError |
| 104 | # --------------------------------------------------------------------------- |
| 105 | |
| 106 | def test_wrong_compression_flag_raises(): |
| 107 | from muse.core.mpack import WireFrameError |
| 108 | from musehub.services.musehub_wire import iter_wire_frames_grpc |
| 109 | |
| 110 | mwp = _make_mwp_frame("E", _pack({"t": "E"})) |
| 111 | bad_stream = b"\x01" + struct.pack(">I", len(mwp)) + mwp # flag=1 = compressed |
| 112 | |
| 113 | async def _run_bad(): |
| 114 | async def _iter(): |
| 115 | yield bad_stream |
| 116 | async for _ in iter_wire_frames_grpc(_iter()): |
| 117 | pass |
| 118 | |
| 119 | with pytest.raises(WireFrameError, match="compress"): |
| 120 | _run(_run_bad()) |
| 121 | |
| 122 | |
| 123 | # --------------------------------------------------------------------------- |
| 124 | # E. Truncated gRPC header raises WireFrameError |
| 125 | # --------------------------------------------------------------------------- |
| 126 | |
| 127 | def test_truncated_grpc_header_raises(): |
| 128 | from muse.core.mpack import WireFrameError |
| 129 | from musehub.services.musehub_wire import iter_wire_frames_grpc |
| 130 | |
| 131 | async def _run_bad(): |
| 132 | async def _iter(): |
| 133 | yield b"\x00\x00\x00" # only 3 bytes — header needs 5 |
| 134 | |
| 135 | async for _ in iter_wire_frames_grpc(_iter()): |
| 136 | pass |
| 137 | |
| 138 | with pytest.raises(WireFrameError): |
| 139 | _run(_run_bad()) |
| 140 | |
| 141 | |
| 142 | # --------------------------------------------------------------------------- |
| 143 | # F. Truncated gRPC message body raises WireFrameError |
| 144 | # --------------------------------------------------------------------------- |
| 145 | |
| 146 | def test_truncated_grpc_body_raises(): |
| 147 | from muse.core.mpack import WireFrameError |
| 148 | from musehub.services.musehub_wire import iter_wire_frames_grpc |
| 149 | |
| 150 | async def _run_bad(): |
| 151 | async def _iter(): |
| 152 | # says 100 bytes but only provides 10 |
| 153 | yield b"\x00" + struct.pack(">I", 100) + b"x" * 10 |
| 154 | |
| 155 | async for _ in iter_wire_frames_grpc(_iter()): |
| 156 | pass |
| 157 | |
| 158 | with pytest.raises(WireFrameError): |
| 159 | _run(_run_bad()) |
| 160 | |
| 161 | |
| 162 | # --------------------------------------------------------------------------- |
| 163 | # G. Full roundtrip: client grpc_frame(MuseWireFrameWriter.wrap) → server iter |
| 164 | # --------------------------------------------------------------------------- |
| 165 | |
| 166 | def test_full_client_server_roundtrip(): |
| 167 | from muse.core.mpack import grpc_frame, MuseWireFrameWriter |
| 168 | |
| 169 | fw = MuseWireFrameWriter() |
| 170 | payload = _pack({"t": "E"}) |
| 171 | mwp = fw.wrap(frame_type="E", payload=payload) |
| 172 | stream = grpc_frame(mwp) |
| 173 | |
| 174 | results = _run(_collect(stream)) |
| 175 | assert len(results) == 1 |
| 176 | envelope, decoded_payload = results[0] |
| 177 | assert envelope["ft"] == "E" |
| 178 | assert decoded_payload == payload |
File History
1 commit
sha256:9590cee1e0ccd6c76528f005b95d634d80f5019f0dcb7c371e149adc31d1fb65
refactor: enforce gRPC framing on all MWP wire traffic
Sonnet 4.6
minor
⚠
155 days ago