chore: Update test coverage for protocol parsing

Author: allenporter

roborock/protocols/b01_q10_protocol.py

@@ -36,9 +36,11 @@ def _convert_datapoints(datapoints: dict[str, Any], message: RoborockMessage) ->
     """Convert the 'dps' dictionary keys from strings to B01_Q10_DP enums."""
     result = {}
     for key, value in datapoints.items():
-        if not isinstance(key, str):
-            raise RoborockException(f"Invalid B01 message format: 'dps' keys should be strings for {message.payload!r}")
-        dps = B01_Q10_DP.from_code(int(key))
+        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
+        dps = B01_Q10_DP.from_code(code)
         result[dps] = value
     return result
 
@@ -49,16 +51,29 @@ def decode_rpc_response(message: RoborockMessage) -> dict[B01_Q10_DP, Any]:
         raise RoborockException("Invalid B01 message format: missing payload")
     try:
         payload = json.loads(message.payload.decode())
-    except (json.JSONDecodeError, TypeError, UnicodeDecodeError) as e:
-        raise RoborockException(f"Invalid B01 message payload: {e} for {message.payload!r}") from e
+    except (json.JSONDecodeError, UnicodeDecodeError) as e:
+        raise RoborockException(f"Invalid B01 json payload: {e} for {message.payload!r}") from e
 
-    datapoints = payload.get("dps", {})
+    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}")
 
-    result = _convert_datapoints(datapoints, message)
-    # The COMMON response contains nested datapoints that also need conversion
+    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 that also need conversion.
+    # We will parse that here for now, but may move elsewhere as we add more
+    # complex response parsing.
     if common_result := result.get(B01_Q10_DP.COMMON):
-        common_dps_result = _convert_datapoints(common_result, message)
+        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[B01_Q10_DP.COMMON] = common_dps_result
+
     return result

tests/protocols/__snapshots__/test_b01_q10_protocol.ambr

@@ -10,7 +10,7 @@
   '''
   {
     "dpCommon": {
-      "90": 0
+      "dpFault": 0
     }
   }
   '''
@@ -19,52 +19,52 @@
   '''
   {
     "dpCommon": {
-      "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": {
+      "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
       },
-      "80": 0,
-      "81": {
+      "dpAreaUnit": 0,
+      "dpNetInfo": {
         "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": {
+      "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
       },
-      "93": 1,
-      "96": 0
+      "dpTimerType": 1,
+      "dpAddCleanState": 0
     },
     "dpStatus": 8,
     "dpBattery": 100,

tests/protocols/test_b01_q10_protocol.py

@@ -6,12 +6,11 @@
 from typing import Any
 
 import pytest
-from Crypto.Cipher import AES
-from Crypto.Util.Padding import unpad
 from freezegun import freeze_time
 from syrupy import SnapshotAssertion
 
 from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP
+from roborock.exceptions import RoborockException
 from roborock.protocols.b01_q10_protocol import (
     decode_rpc_response,
     encode_mqtt_payload,
@@ -49,6 +48,32 @@ def test_decode_rpc_payload(filename: str, snapshot: SnapshotAssertion) -> None:
     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)
+
+
 @pytest.mark.parametrize(
     ("command", "params"),
     [
@@ -63,7 +88,6 @@ def test_encode_mqtt_payload(command: B01_Q10_DP, params: dict[str, Any]) -> Non
     assert message.protocol == RoborockMessageProtocol.RPC_REQUEST
     assert message.version == b"B01"
     assert message.payload is not None
-    unpadded = unpad(message.payload, AES.block_size)
-    decoded_json = json.loads(unpadded.decode("utf-8"))
+    decoded_json = json.loads(message.payload.decode("utf-8"))
 
     assert decoded_json == {"dps": {"102": {}}}