diff --git a/roborock/data/code_mappings.py b/roborock/data/code_mappings.py index a0991608..e9066f1f 100644 --- a/roborock/data/code_mappings.py +++ b/roborock/data/code_mappings.py @@ -55,6 +55,7 @@ class RoborockModeEnum(StrEnum): """A custom StrEnum that also stores an integer code for each member.""" code: int + """The integer code associated with the enum member.""" def __new__(cls, value: str, code: int) -> RoborockModeEnum: """Creates a new enum member.""" @@ -68,7 +69,18 @@ def from_code(cls, code: int) -> RoborockModeEnum: for member in cls: if member.code == code: return member - raise ValueError(f"{code} is not a valid code for {cls.__name__}") + message = f"{code} is not a valid code for {cls.__name__}" + if message not in completed_warnings: + completed_warnings.add(message) + _LOGGER.warning(message) + raise ValueError(message) + + @classmethod + def from_code_optional(cls, code: int) -> RoborockModeEnum | None: + try: + return cls.from_code(code) + except ValueError: + return None @classmethod def from_value(cls, value: str) -> RoborockModeEnum: diff --git a/roborock/protocols/b01_q10_protocol.py b/roborock/protocols/b01_q10_protocol.py new file mode 100644 index 00000000..06d57921 --- /dev/null +++ b/roborock/protocols/b01_q10_protocol.py @@ -0,0 +1,88 @@ +"""Roborock B01 Protocol encoding and decoding.""" + +import json +import logging +from typing import Any + +from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP +from roborock.exceptions import RoborockException +from roborock.roborock_message import ( + RoborockMessage, + RoborockMessageProtocol, +) + +_LOGGER = logging.getLogger(__name__) + +B01_VERSION = b"B01" +ParamsType = list | dict | int | None + + +def encode_mqtt_payload(command: B01_Q10_DP, params: ParamsType) -> RoborockMessage: + """Encode payload for B01 Q10 commands over MQTT. + + This does not perform any special encoding for the command parameters and expects + them to already be in a request specific format. + """ + dps_data = { + "dps": { + # Important: some commands use falsy values so only default to `{}` when params is actually None. + command.code: params if params is not None else {}, + } + } + return RoborockMessage( + protocol=RoborockMessageProtocol.RPC_REQUEST, + version=B01_VERSION, + payload=json.dumps(dps_data).encode("utf-8"), + ) + + +def _convert_datapoints(datapoints: dict[str, Any], message: RoborockMessage) -> dict[B01_Q10_DP, Any]: + """Convert the 'dps' dictionary keys from strings to B01_Q10_DP enums.""" + result: dict[B01_Q10_DP, Any] = {} + for key, value in datapoints.items(): + try: + code = int(key) + except ValueError as e: + raise ValueError(f"dps key is not a valid integer: {e} for {message.payload!r}") from e + if (dps := B01_Q10_DP.from_code_optional(code)) is not None: + # Update from_code to use `Self` on newer python version to remove this type ignore + result[dps] = value # type: ignore[index] + return result + + +def decode_rpc_response(message: RoborockMessage) -> dict[B01_Q10_DP, Any]: + """Decode a B01 Q10 RPC_RESPONSE message. + + This does not perform any special decoding for the response body, but does + convert the 'dps' keys from strings to B01_Q10_DP enums. + """ + if not message.payload: + raise RoborockException("Invalid B01 message format: missing payload") + try: + payload = json.loads(message.payload.decode()) + except (json.JSONDecodeError, UnicodeDecodeError) as e: + raise RoborockException(f"Invalid B01 json payload: {e} for {message.payload!r}") from e + + if (datapoints := payload.get("dps")) is None: + raise RoborockException(f"Invalid B01 json payload: missing 'dps' for {message.payload!r}") + if not isinstance(datapoints, dict): + raise RoborockException(f"Invalid B01 message format: 'dps' should be a dictionary for {message.payload!r}") + + try: + result = _convert_datapoints(datapoints, message) + except ValueError as e: + raise RoborockException(f"Invalid B01 message format: {e}") from e + + # The COMMON response contains nested datapoints need conversion. To simplify + # response handling at higher levels we flatten these into the main result. + if B01_Q10_DP.COMMON in result: + common_result = result.pop(B01_Q10_DP.COMMON) + if not isinstance(common_result, dict): + raise RoborockException(f"Invalid dpCommon format: expected dict, got {type(common_result).__name__}") + try: + common_dps_result = _convert_datapoints(common_result, message) + except ValueError as e: + raise RoborockException(f"Invalid dpCommon format: {e}") from e + result.update(common_dps_result) + + return result diff --git a/tests/data/test_code_mappings.py b/tests/data/test_code_mappings.py index cf787148..c59460ac 100644 --- a/tests/data/test_code_mappings.py +++ b/tests/data/test_code_mappings.py @@ -20,6 +20,11 @@ def test_invalid_from_code() -> None: B01_Q10_DP.from_code(999999) +def test_invalid_from_code_optional() -> None: + """Test invalid from_code_optional method.""" + assert B01_Q10_DP.from_code_optional(999999) is None + + def test_from_name() -> None: """Test from_name method.""" assert B01_Q10_DP.START_CLEAN == B01_Q10_DP.from_name("START_CLEAN") diff --git a/tests/protocols/__snapshots__/test_b01_q10_protocol.ambr b/tests/protocols/__snapshots__/test_b01_q10_protocol.ambr new file mode 100644 index 00000000..0eae2538 --- /dev/null +++ b/tests/protocols/__snapshots__/test_b01_q10_protocol.ambr @@ -0,0 +1,91 @@ +# serializer version: 1 +# name: test_decode_rpc_payload[dpBattery] + ''' + { + "dpBattery": 100 + } + ''' +# --- +# name: test_decode_rpc_payload[dpRequetdps] + ''' + { + "dpStatus": 8, + "dpBattery": 100, + "dpfunLevel": 2, + "dpWaterLevel": 1, + "dpMainBrushLife": 0, + "dpSideBrushLife": 0, + "dpFilterLife": 0, + "dpCleanCount": 1, + "dpCleanMode": 1, + "dpCleanTaskType": 0, + "dpBackType": 5, + "dpBreakpointClean": 0, + "dpValleyPointCharging": false, + "dpRobotCountryCode": "us", + "dpUserPlan": 0, + "dpNotDisturb": 1, + "dpVolume": 74, + "dpTotalCleanArea": 0, + "dpTotalCleanCount": 0, + "dpTotalCleanTime": 0, + "dpDustSwitch": 1, + "dpMopState": 1, + "dpAutoBoost": 0, + "dpChildLock": 0, + "dpDustSetting": 0, + "dpMapSaveSwitch": true, + "dpRecendCleanRecord": false, + "dpCleanTime": 0, + "dpMultiMapSwitch": 1, + "dpSensorLife": 0, + "dpCleanArea": 0, + "dpCarpetCleanType": 0, + "dpCleanLine": 0, + "dpTimeZone": { + "timeZoneCity": "America/Los_Angeles", + "timeZoneSec": -28800 + }, + "dpAreaUnit": 0, + "dpNetInfo": { + "ipAdress": "1.1.1.2", + "mac": "99:AA:88:BB:77:CC", + "signal": -50, + "wifiName": "wifi-network-name" + }, + "dpRobotType": 1, + "dpLineLaserObstacleAvoidance": 1, + "dpCleanProgess": 100, + "dpGroundClean": 0, + "dpFault": 0, + "dpNotDisturbExpand": { + "disturb_dust_enable": 1, + "disturb_light": 1, + "disturb_resume_clean": 1, + "disturb_voice": 1 + }, + "dpTimerType": 1, + "dpAddCleanState": 0 + } + ''' +# --- +# name: test_decode_rpc_payload[dpStatus-dpCleanTaskType] + ''' + { + "dpStatus": 8, + "dpCleanTaskType": 0 + } + ''' +# --- +# name: test_encode_mqtt_payload[dpRequetdps-None] + b'{"dps": {"102": {}}}' +# --- +# name: test_encode_mqtt_payload[dpRequetdps-params0] + b'{"dps": {"102": {}}}' +# --- +# name: test_encode_mqtt_payload[dpStartClean-params2] + b'{"dps": {"201": {"cmd": 1}}}' +# --- +# name: test_encode_mqtt_payload[dpWaterLevel-2] + b'{"dps": {"124": 2}}' +# --- diff --git a/tests/protocols/test_b01_q10_protocol.py b/tests/protocols/test_b01_q10_protocol.py new file mode 100644 index 00000000..fa373f38 --- /dev/null +++ b/tests/protocols/test_b01_q10_protocol.py @@ -0,0 +1,116 @@ +"""Tests for the B01 protocol message encoding and decoding.""" + +import json +import pathlib +from collections.abc import Generator +from typing import Any + +import pytest +from freezegun import freeze_time +from syrupy import SnapshotAssertion + +from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP, YXWaterLevel +from roborock.exceptions import RoborockException +from roborock.protocols.b01_q10_protocol import ( + decode_rpc_response, + encode_mqtt_payload, +) +from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol + +TESTDATA_PATH = pathlib.Path("tests/protocols/testdata/b01_q10_protocol/") +TESTDATA_FILES = list(TESTDATA_PATH.glob("*.json")) +TESTDATA_IDS = [x.stem for x in TESTDATA_FILES] + + +@pytest.fixture(autouse=True) +def fixed_time_fixture() -> Generator[None, None, None]: + """Fixture to freeze time for predictable request IDs.""" + with freeze_time("2025-01-20T12:00:00"): + yield + + +@pytest.mark.parametrize("filename", TESTDATA_FILES, ids=TESTDATA_IDS) +def test_decode_rpc_payload(filename: str, snapshot: SnapshotAssertion) -> None: + """Test decoding a B01 RPC response protocol message.""" + with open(filename, "rb") as f: + payload = f.read() + + message = RoborockMessage( + protocol=RoborockMessageProtocol.RPC_RESPONSE, + payload=payload, + seq=12750, + version=b"B01", + random=97431, + timestamp=1652547161, + ) + + decoded_message = decode_rpc_response(message) + assert json.dumps(decoded_message, indent=2) == snapshot + + +@pytest.mark.parametrize( + ("payload", "expected_error_message"), + [ + (b"", "missing payload"), + (b"n", "Invalid B01 json payload"), + (b"{}", "missing 'dps'"), + (b'{"dps": []}', "'dps' should be a dictionary"), + (b'{"dps": {"not_a_number": 123}}', "dps key is not a valid integer"), + (b'{"dps": {"101": 123}}', "Invalid dpCommon format: expected dict"), + (b'{"dps": {"101": {"not_a_number": 123}}}', "Invalid dpCommon format: dps key is not a valid intege"), + ], +) +def test_decode_invalid_rpc_payload(payload: bytes, expected_error_message: str) -> None: + """Test decoding a B01 RPC response protocol message.""" + message = RoborockMessage( + protocol=RoborockMessageProtocol.RPC_RESPONSE, + payload=payload, + seq=12750, + version=b"B01", + random=97431, + timestamp=1652547161, + ) + with pytest.raises(RoborockException, match=expected_error_message): + decode_rpc_response(message) + + +def test_decode_unknown_dps_code() -> None: + """Test decoding a B01 RPC response protocol message.""" + message = RoborockMessage( + protocol=RoborockMessageProtocol.RPC_RESPONSE, + payload=b'{"dps": {"909090": 123, "122":100}}', + seq=12750, + version=b"B01", + random=97431, + timestamp=1652547161, + ) + + decoded_message = decode_rpc_response(message) + assert decoded_message == { + B01_Q10_DP.BATTERY: 100, + } + + +@pytest.mark.parametrize( + ("command", "params"), + [ + (B01_Q10_DP.REQUETDPS, {}), + (B01_Q10_DP.REQUETDPS, None), + (B01_Q10_DP.START_CLEAN, {"cmd": 1}), + (B01_Q10_DP.WATER_LEVEL, YXWaterLevel.MIDDLE.code), + ], +) +def test_encode_mqtt_payload(command: B01_Q10_DP, params: dict[str, Any], snapshot) -> None: + """Test encoding of MQTT payload for B01 Q10 commands.""" + + message = encode_mqtt_payload(command, params) + assert isinstance(message, RoborockMessage) + assert message.protocol == RoborockMessageProtocol.RPC_REQUEST + assert message.version == b"B01" + assert message.payload is not None + + # Snapshot the raw payload to ensure stable encoding. We verify it is + # valid json + assert snapshot == message.payload + + json.loads(message.payload.decode()) diff --git a/tests/protocols/testdata/b01_q10_protocol/dpBattery.json b/tests/protocols/testdata/b01_q10_protocol/dpBattery.json new file mode 100644 index 00000000..b7a76fda --- /dev/null +++ b/tests/protocols/testdata/b01_q10_protocol/dpBattery.json @@ -0,0 +1 @@ +{"dps":{"122":100},"t":1766800902} diff --git a/tests/protocols/testdata/b01_q10_protocol/dpRequetdps.json b/tests/protocols/testdata/b01_q10_protocol/dpRequetdps.json new file mode 100644 index 00000000..13d52243 --- /dev/null +++ b/tests/protocols/testdata/b01_q10_protocol/dpRequetdps.json @@ -0,0 +1 @@ +{"dps":{"101":{"104":0,"105":false,"109":"us","207":0,"25":1,"26":74,"29":0,"30":0,"31":0,"37":1,"40":1,"45":0,"47":0,"50":0,"51":true,"53":false,"6":0,"60":1,"67":0,"7":0,"76":0,"78":0,"79":{"timeZoneCity":"America/Los_Angeles","timeZoneSec":-28800},"80":0,"81":{"ipAdress":"1.1.1.2","mac":"99:AA:88:BB:77:CC","signal":-50,"wifiName":"wifi-network-name"},"83":1,"86":1,"87":100,"88":0,"90":0,"92":{"disturb_dust_enable":1,"disturb_light":1,"disturb_resume_clean":1,"disturb_voice":1},"93":1,"96":0},"121":8,"122":100,"123":2,"124":1,"125":0,"126":0,"127":0,"136":1,"137":1,"138":0,"139":5},"t":1766802312} diff --git a/tests/protocols/testdata/b01_q10_protocol/dpStatus-dpCleanTaskType.json b/tests/protocols/testdata/b01_q10_protocol/dpStatus-dpCleanTaskType.json new file mode 100644 index 00000000..ed9de954 --- /dev/null +++ b/tests/protocols/testdata/b01_q10_protocol/dpStatus-dpCleanTaskType.json @@ -0,0 +1 @@ +{"dps":{"121":8,"138":0},"t":1766800904}