chore: tests

Author: Lash-L

roborock/protocols/b01_protocol.py

@@ -47,8 +47,9 @@ def decode_rpc_response(message: RoborockMessage) -> dict[int, Any]:
         raise RoborockException("Invalid B01 message format: missing payload")
     try:
         unpadded = unpad(message.payload, AES.block_size)
-    except ValueError as err:
-        raise RoborockException(f"Unable to unpad B01 payload: {err}")
+    except ValueError:
+        # It would be better to fail down the line.
+        unpadded = message.payload
 
     try:
         payload = json.loads(unpadded.decode())

tests/devices/traits/b01/test_init.py

@@ -0,0 +1,130 @@
+import json
+from typing import Any
+from unittest.mock import patch
+
+import pytest
+from Crypto.Cipher import AES
+from Crypto.Util.Padding import pad, unpad
+
+from roborock.data.b01_q7 import WorkStatusMapping
+from roborock.devices.traits.b01.q7 import Q7PropertiesApi
+from roborock.protocols.b01_protocol import B01_VERSION
+from roborock.roborock_message import RoborockB01Props, RoborockMessage, RoborockMessageProtocol
+from tests.conftest import FakeChannel
+
+
+def build_b01_message(message: dict[Any, Any], msg_id: str = "123456789", seq: int = 2020) -> RoborockMessage:
+    """Build an encoded B01 RPC response message."""
+    dps_payload = {
+        "dps": {
+            "10000": json.dumps(
+                {
+                    "msgId": msg_id,
+                    "data": message,
+                }
+            )
+        }
+    }
+    return RoborockMessage(
+        protocol=RoborockMessageProtocol.RPC_RESPONSE,
+        payload=pad(
+            json.dumps(dps_payload).encode(),
+            AES.block_size,
+        ),
+        version=b"B01",
+        seq=seq,
+    )
+
+
+@pytest.fixture(name="fake_channel")
+def fake_channel_fixture() -> FakeChannel:
+    return FakeChannel()
+
+
+@pytest.fixture(name="q7_api")
+def q7_api_fixture(fake_channel: FakeChannel) -> Q7PropertiesApi:
+    return Q7PropertiesApi(fake_channel)  # type: ignore[arg-type]
+
+
+async def test_q7_api_query_values(q7_api: Q7PropertiesApi, fake_channel: FakeChannel):
+    """Test that Q7PropertiesApi correctly converts raw values."""
+    expected_msg_id = "123456789"
+
+    # We need to construct the expected result based on the mappings
+    # status: 1 -> WAITING_FOR_ORDERS
+    # wind: 1 -> STANDARD
+    response_data = {
+        "status": 1,
+        "wind": 1,
+        "battery": 100,
+    }
+
+    # Patch get_next_int to return our expected msg_id so the channel waits for it
+    with patch("roborock.devices.b01_channel.get_next_int", return_value=int(expected_msg_id)):
+        # Queue the response
+        fake_channel.response_queue.append(build_b01_message(response_data, msg_id=expected_msg_id))
+
+        result = await q7_api.query_values(
+            [
+                RoborockB01Props.STATUS,
+                RoborockB01Props.WIND,
+            ]
+        )
+
+    assert result is not None
+    assert result.status == WorkStatusMapping.WAITING_FOR_ORDERS
+    # wind might be mapped to SCWindMapping.STANDARD (1)
+    # let's verify checking the prop definition in B01Props
+    # wind: SCWindMapping | None = None
+    # SCWindMapping.STANDARD is 1 ('balanced')
+    from roborock.data.b01_q7 import SCWindMapping
+
+    assert result.wind == SCWindMapping.STANDARD
+
+    assert len(fake_channel.published_messages) == 1
+    message = fake_channel.published_messages[0]
+    assert message.protocol == RoborockMessageProtocol.RPC_REQUEST
+    assert message.version == B01_VERSION
+
+    # Verify request payload
+    payload_data = json.loads(unpad(message.payload, AES.block_size))
+    # {"dps": {"10000": {"method": "prop.get", "msgId": "123456789", "params": {"property": ["status", "wind"]}}}}
+    assert "dps" in payload_data
+    assert "10000" in payload_data["dps"]
+    inner = payload_data["dps"]["10000"]
+    assert inner["method"] == "prop.get"
+    assert inner["msgId"] == expected_msg_id
+    assert inner["params"] == {"property": [RoborockB01Props.STATUS, RoborockB01Props.WIND]}
+
+
+@pytest.mark.parametrize(
+    ("query", "response_data", "expected_status"),
+    [
+        (
+            [RoborockB01Props.STATUS],
+            {"status": 2},
+            WorkStatusMapping.PAUSED,
+        ),
+        (
+            [RoborockB01Props.STATUS],
+            {"status": 5},
+            WorkStatusMapping.SWEEP_MOPING,
+        ),
+    ],
+)
+async def test_q7_response_value_mapping(
+    query: list[RoborockB01Props],
+    response_data: dict[str, Any],
+    expected_status: WorkStatusMapping,
+    q7_api: Q7PropertiesApi,
+    fake_channel: FakeChannel,
+):
+    """Test Q7PropertiesApi value mapping for different statuses."""
+    msg_id = "987654321"
+
+    with patch("roborock.devices.b01_channel.get_next_int", return_value=int(msg_id)):
+        fake_channel.response_queue.append(build_b01_message(response_data, msg_id=msg_id))
+
+        result = await q7_api.query_values(query)
+
+    assert result is not None

tests/protocols/test_b01_q07_protocol.py

@@ -0,0 +1,74 @@
+"""Tests for the B01 protocol message encoding and decoding."""
+
+import json
+import pathlib
+from collections.abc import Generator
+
+import pytest
+from Crypto.Cipher import AES
+from Crypto.Util.Padding import unpad
+from freezegun import freeze_time
+from syrupy import SnapshotAssertion
+
+from roborock.protocols.b01_protocol import (
+    decode_rpc_response,
+    encode_mqtt_payload,
+)
+from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
+
+TESTDATA_PATH = pathlib.Path("tests/protocols/testdata/b01_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(
+    ("dps", "command", "params", "msg_id"),
+    [
+        (
+            10000,
+            "prop.get",
+            {"property": ["status", "fault"]},
+            "123456789",
+        ),
+    ],
+)
+def test_encode_mqtt_payload(dps: int, command: str, params: dict[str, list[str]], msg_id: str) -> None:
+    """Test encoding of MQTT payload for B01 commands."""
+
+    message = encode_mqtt_payload(dps, command, params, msg_id)
+    assert isinstance(message, RoborockMessage)
+    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"))
+
+    assert decoded_json["dps"][str(dps)]["method"] == command
+    assert decoded_json["dps"][str(dps)]["msgId"] == msg_id
+    assert decoded_json["dps"][str(dps)]["params"] == params

tests/protocols/testdata/b01_protocol/q7/get_prop.json

@@ -0,0 +1 @@
+{"t":1765660648,"dps":{"10001":"{\"msgId\":\"200000000001\",\"code\":0,\"method\":\"prop.get\",\"data\":{\"status\":4,\"main_brush\":4088}}"}}