Add property testing for ser-de (#27)

Author: vrslev

.gitignore

@@ -1,4 +1,5 @@
 .coverage
+.hypothesis
 .venv
 dist
 uv.lock

pyproject.toml

@@ -25,8 +25,9 @@ dev-dependencies = [
     "mypy~=1.10.0",
     "pytest-cov~=5.0.0",
     "pytest~=8.2.2",
-    "ruff~=0.4.8",
+    "ruff~=0.4.9",
     "uvloop~=0.19.0",
+    "hypothesis~=6.103.2",
 ]
 
 [build-system]

stompman/serde.py

@@ -1,6 +1,7 @@
 import struct
 from collections import deque
 from collections.abc import Iterator
+from contextlib import suppress
 from dataclasses import dataclass, field
 from typing import Any, Final, cast
 
@@ -32,13 +33,12 @@
     "\n": "\\n",
     ":": "\\c",
     "\\": "\\\\",
-    "\r": "\\r",
+    "\r": "",  # [\r]\n is newline, therefore can't be used in header
 }
 HEADER_UNESCAPE_CHARS: Final = {
-    b"n": b"\n",
+    b"n": NEWLINE,
     b"c": b":",
     b"\\": b"\\",
-    b"r": b"\r",
 }
 
 
@@ -89,51 +89,62 @@ def dump_frame(frame: AnyClientFrame | AnyServerFrame) -> bytes:
 
 
 def unescape_byte(byte: bytes, previous_byte: bytes | None) -> bytes | None:
+    if previous_byte == b"\\":
+        return HEADER_UNESCAPE_CHARS.get(byte)
     if byte == b"\\":
         return None
-
-    if previous_byte == b"\\":
-        return HEADER_UNESCAPE_CHARS.get(byte, byte)
-
     return byte
 
 
-def parse_headers(buffer: list[bytes]) -> tuple[str, str] | None:
+def parse_header(buffer: list[bytes]) -> tuple[str, str] | None:
     key_buffer: list[bytes] = []
     key_parsed = False
     value_buffer: list[bytes] = []
 
     previous_byte = None
+    just_escaped_line = False
+
     for byte in buffer:
         if byte == b":":
             if key_parsed:
                 return None
             key_parsed = True
-
-        elif (unescaped_byte := unescape_byte(byte, previous_byte)) is not None:
+        elif just_escaped_line:
+            just_escaped_line = False
+            if byte != b"\\":
+                (value_buffer if key_parsed else key_buffer).append(byte)
+        elif unescaped_byte := unescape_byte(byte, previous_byte):
+            just_escaped_line = True
             (value_buffer if key_parsed else key_buffer).append(unescaped_byte)
 
         previous_byte = byte
 
-    return (b"".join(key_buffer).decode(), b"".join(value_buffer).decode()) if key_parsed else None
+    if key_parsed:
+        with suppress(UnicodeDecodeError):
+            return (b"".join(key_buffer).decode(), b"".join(value_buffer).decode())
+
+    return None
+
+
+def make_frame_from_parts(command: bytes, headers: dict[str, str], body: bytes) -> AnyClientFrame | AnyServerFrame:
+    frame_type = COMMANDS_TO_FRAMES[command]
+    return (
+        frame_type(headers=cast(Any, headers), body=body)  # type: ignore[call-arg]
+        if frame_type in FRAMES_WITH_BODY
+        else frame_type(headers=cast(Any, headers))  # type: ignore[call-arg]
+    )
 
 
 def parse_lines_into_frame(lines: deque[list[bytes]]) -> AnyClientFrame | AnyServerFrame:
     command = b"".join(lines.popleft())
     headers = {}
 
     while line := lines.popleft():
-        header = parse_headers(line)
+        header = parse_header(line)
         if header and header[0] not in headers:
             headers[header[0]] = header[1]
     body = b"".join(lines.popleft()) if lines else b""
-
-    frame_type = COMMANDS_TO_FRAMES[command]
-    return (
-        frame_type(headers=cast(Any, headers), body=body)  # type: ignore[call-arg]
-        if frame_type in FRAMES_WITH_BODY
-        else frame_type(headers=cast(Any, headers))  # type: ignore[call-arg]
-    )
+    return make_frame_from_parts(command=command, headers=headers, body=body)
 
 
 @dataclass
@@ -149,10 +160,7 @@ def _reset(self) -> None:
         self._current_line = []
 
     def parse_frames_from_chunk(self, chunk: bytes) -> Iterator[AnyClientFrame | AnyServerFrame | HeartbeatFrame]:
-        buffer = deque(iter_bytes(chunk))
-        while buffer:
-            byte = buffer.popleft()
-
+        for byte in iter_bytes(chunk):
             if byte == NULL:
                 if self._headers_processed:
                     self._lines.append(self._current_line)

tests/integration.py

@@ -1,13 +1,26 @@
 import asyncio
 import os
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Callable
 from contextlib import asynccontextmanager
+from itertools import starmap
 from uuid import uuid4
 
 import pytest
+from hypothesis import given, strategies
 
 import stompman
-from stompman.errors import ConnectionLostError
+from stompman import AnyClientFrame, AnyServerFrame, ConnectionLostError, HeartbeatFrame
+from stompman.serde import (
+    COMMANDS_TO_FRAMES,
+    NEWLINE,
+    NULL,
+    FrameParser,
+    dump_frame,
+    dump_header,
+    iter_bytes,
+    make_frame_from_parts,
+    parse_header,
+)
 
 pytestmark = pytest.mark.anyio
 
@@ -113,3 +126,55 @@ async def test_raises_connection_lost_error_in_listen(client: stompman.Client) -
     client.read_timeout = 0
     with pytest.raises(ConnectionLostError):
         [event async for event in client.listen()]
+
+
+def generate_frames(
+    cases: list[tuple[bytes, list[AnyClientFrame | AnyServerFrame | HeartbeatFrame]]],
+) -> tuple[list[bytes], list[AnyClientFrame | AnyServerFrame | HeartbeatFrame]]:
+    all_bytes, all_frames = [], []
+
+    for noise, frames in cases:
+        current_all_bytes = []
+        if noise:
+            current_all_bytes.append(noise + NEWLINE)
+
+        for frame in frames:
+            current_all_bytes.append(NEWLINE if isinstance(frame, HeartbeatFrame) else dump_frame(frame))
+            all_frames.append(frame)
+
+        all_bytes.append(b"".join(current_all_bytes))
+
+    return all_bytes, all_frames
+
+
+def bytes_not_contains(*avoided: bytes) -> Callable[[bytes], bool]:
+    return lambda checked: all(item not in checked for item in avoided)
+
+
+noise_bytes_strategy = strategies.binary().filter(bytes_not_contains(NEWLINE, NULL))
+header_value_strategy = strategies.text().filter(lambda text: "\x00" not in text)
+headers_strategy = strategies.dictionaries(header_value_strategy, header_value_strategy).map(
+    lambda headers: dict(
+        parsed_header
+        for header in starmap(dump_header, headers.items())
+        if (parsed_header := parse_header(list(iter_bytes(header))))
+    )
+)
+frame_strategy = strategies.just(HeartbeatFrame()) | strategies.builds(
+    make_frame_from_parts,
+    command=strategies.sampled_from(tuple(COMMANDS_TO_FRAMES.keys())),
+    headers=headers_strategy,
+    body=strategies.binary().filter(bytes_not_contains(NULL)),
+)
+
+
+@given(
+    strategies.builds(
+        generate_frames,
+        strategies.lists(strategies.tuples(noise_bytes_strategy, strategies.lists(frame_strategy))),
+    ),
+)
+def test_parsing(case: tuple[list[bytes], list[AnyClientFrame | AnyServerFrame | HeartbeatFrame]]) -> None:
+    stream_chunks, expected_frames = case
+    parser = FrameParser()
+    assert [frame for chunk in stream_chunks for frame in parser.parse_frames_from_chunk(chunk)] == expected_frames

tests/test_client.py

@@ -20,10 +20,14 @@
     ConnectedFrame,
     ConnectFrame,
     ConnectionConfirmationTimeoutError,
+    ConnectionParameters,
     DisconnectFrame,
+    ErrorEvent,
     ErrorFrame,
     FailedAllConnectAttemptsError,
+    HeartbeatEvent,
     HeartbeatFrame,
+    MessageEvent,
     MessageFrame,
     NackFrame,
     ReceiptFrame,
@@ -32,7 +36,6 @@
     UnsubscribeFrame,
     UnsupportedProtocolVersionError,
 )
-from stompman.client import ConnectionParameters, ErrorEvent, HeartbeatEvent, MessageEvent
 from stompman.errors import ConnectionLostError
 
 pytestmark = pytest.mark.anyio

tests/test_connection.py

@@ -7,8 +7,16 @@
 
 import pytest
 
-from stompman import AnyServerFrame, ConnectedFrame, Connection, ConnectionLostError, HeartbeatFrame
-from stompman.frames import BeginFrame, CommitFrame
+from stompman import (
+    AnyServerFrame,
+    BeginFrame,
+    CommitFrame,
+    ConnectedFrame,
+    Connection,
+    ConnectionLostError,
+    HeartbeatFrame,
+)
+from stompman.serde import NEWLINE
 
 pytestmark = pytest.mark.anyio
 
@@ -95,7 +103,7 @@ async def take_frames(count: int) -> list[AnyServerFrame]:
     MockWriter.wait_closed.assert_called_once_with()
     MockWriter.drain.assert_called_once_with()
     MockReader.read.mock_calls = [mock.call(max_chunk_size)] * len(read_bytes)  # type: ignore[assignment]
-    assert MockWriter.write.mock_calls == [mock.call(b"\n"), mock.call(b"COMMIT\ntransaction:transaction\n\n\x00")]
+    assert MockWriter.write.mock_calls == [mock.call(NEWLINE), mock.call(b"COMMIT\ntransaction:transaction\n\n\x00")]
 
 
 async def test_connection_close_connection_error(monkeypatch: pytest.MonkeyPatch) -> None:

tests/test_frame_serde.py

@@ -1,14 +1,16 @@
 import pytest
 
 from stompman import (
+    AckFrame,
+    AnyClientFrame,
+    AnyServerFrame,
     ConnectedFrame,
     ConnectFrame,
     ErrorFrame,
     HeartbeatFrame,
     MessageFrame,
 )
-from stompman.frames import AckFrame, AnyClientFrame, AnyServerFrame
-from stompman.serde import FrameParser, dump_frame
+from stompman.serde import NEWLINE, FrameParser, dump_frame
 
 
 @pytest.mark.parametrize(
@@ -18,15 +20,12 @@
         (ConnectedFrame(headers={"version": "1.1"}), (b"CONNECTED\nversion:1.1\n\n\x00")),
         (
             MessageFrame(
-                headers={"destination": "me:123", "message-id": "you\nmore\rextra\\here", "subscription": "hi"},
+                headers={"destination": "me:123", "message-id": "you\nmoreextra\\here", "subscription": "hi"},
                 body=b"I Am The Walrus",
             ),
             (
-                b"MESSAGE\n"
-                b"destination:me\\c123\n"
-                b"message-id:you\\nmore\\rextra\\\\here\nsubscription:hi\n\n"
-                b"I Am The Walrus"
-                b"\x00"
+                b"MESSAGE\ndestination:me\\c123\nmessage-id:you\\nmoreextra\\\\here\nsubscription:hi\n\n"
+                b"I Am The Walrus\x00"
             ),
         ),
     ],
@@ -181,7 +180,7 @@ def test_dump_frame(frame: AnyClientFrame, dumped_frame: bytes) -> None:
                 HeartbeatFrame(),
             ],
         ),
-        (b"\n", [HeartbeatFrame()]),
+        (NEWLINE, [HeartbeatFrame()]),
         # Two headers: only first should be accepted
         (
             b"CONNECTED\naccept-version:1.0\naccept-version:1.1\n\n\x00",