From 94c6aa0ff8b7d2fe57f44d1dca23080d3ae0a7d9 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Sat, 11 Jul 2026 17:18:10 -0700 Subject: [PATCH] feat: truncate large payloads in RoborockMessage __repr__ --- roborock/roborock_message.py | 18 +++++++++++++ tests/test_roborock_message.py | 49 +++++++++++++++++++++++++++++++++- 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/roborock/roborock_message.py b/roborock/roborock_message.py index 72564b38..b78ce5ba 100644 --- a/roborock/roborock_message.py +++ b/roborock/roborock_message.py @@ -274,6 +274,9 @@ class RoborockB01Props(StrEnum): ] +MAX_PAYLOAD_REPR_LEN = 1024 + + @dataclass class RoborockMessage: protocol: RoborockMessageProtocol @@ -282,3 +285,18 @@ class RoborockMessage: version: bytes = b"1.0" random: int = field(default_factory=lambda: get_next_int(10000, 99999)) timestamp: int = field(default_factory=lambda: get_timestamp()) + + def __repr__(self) -> str: + payload_repr = "None" + if self.payload is not None: + if isinstance(self.payload, (bytes, bytearray, str)) and len(self.payload) > MAX_PAYLOAD_REPR_LEN: + r = repr(self.payload[:MAX_PAYLOAD_REPR_LEN]) + quote = r[-1] + payload_repr = f"{r[:-1]}...{quote} (length: {len(self.payload)})" + else: + payload_repr = repr(self.payload) + return ( + f"RoborockMessage(protocol={self.protocol}, payload={payload_repr}, " + f"seq={self.seq}, version={self.version!r}, random={self.random}, " + f"timestamp={self.timestamp})" + ) diff --git a/tests/test_roborock_message.py b/tests/test_roborock_message.py index 2f3ccead..da426d17 100644 --- a/tests/test_roborock_message.py +++ b/tests/test_roborock_message.py @@ -2,7 +2,7 @@ from freezegun import freeze_time -from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol +from roborock.roborock_message import MAX_PAYLOAD_REPR_LEN, RoborockMessage, RoborockMessageProtocol def test_roborock_message() -> None: @@ -23,3 +23,50 @@ def test_roborock_message() -> None: assert message1.seq != message2.seq assert message1.random != message2.random assert message1.timestamp > message2.timestamp + + +def test_roborock_message_repr() -> None: + """Test custom __repr__ of RoborockMessage truncates long payloads.""" + # 1. Payload is None + msg_none = RoborockMessage( + protocol=RoborockMessageProtocol.RPC_REQUEST, + payload=None, + seq=123456, + version=b"1.0", + random=12345, + timestamp=1700000000, + ) + assert repr(msg_none) == ( + "RoborockMessage(protocol=101, payload=None, seq=123456, version=b'1.0', random=12345, timestamp=1700000000)" + ) + + # 2. Payload is short (<= MAX_PAYLOAD_REPR_LEN bytes) + short_payload = b"a" * MAX_PAYLOAD_REPR_LEN + msg_short = RoborockMessage( + protocol=RoborockMessageProtocol.RPC_REQUEST, + payload=short_payload, + seq=123456, + version=b"1.0", + random=12345, + timestamp=1700000000, + ) + assert repr(msg_short) == ( + f"RoborockMessage(protocol=101, payload={repr(short_payload)}, seq=123456, " + "version=b'1.0', random=12345, timestamp=1700000000)" + ) + + # 3. Payload is long (> MAX_PAYLOAD_REPR_LEN bytes) + long_payload = b"a" * (MAX_PAYLOAD_REPR_LEN + 10) + msg_long = RoborockMessage( + protocol=RoborockMessageProtocol.RPC_REQUEST, + payload=long_payload, + seq=123456, + version=b"1.0", + random=12345, + timestamp=1700000000, + ) + expected_payload_repr = f"b'{'a' * MAX_PAYLOAD_REPR_LEN}...' (length: {MAX_PAYLOAD_REPR_LEN + 10})" + assert repr(msg_long) == ( + f"RoborockMessage(protocol=101, payload={expected_payload_repr}, seq=123456, " + "version=b'1.0', random=12345, timestamp=1700000000)" + )