🦎 q7: align B01 map helpers with maintainer feedback

Author: arduano

roborock/devices/rpc/b01_q7_channel.py

@@ -34,13 +34,12 @@ async def _send_command(
     request_message: Q7RequestMessage,
     *,
     response_matcher: Callable[[RoborockMessage], _T | None],
-    timeout_error: str,
 ) -> _T:
     """Publish a B01 command and resolve on the first matching response."""
     roborock_message = encode_mqtt_payload(request_message)
     future: asyncio.Future[_T] = asyncio.get_running_loop().create_future()
 
-    def find_response(response_message: RoborockMessage) -> None:
+    def on_message(response_message: RoborockMessage) -> None:
         if future.done():
             return
         try:
@@ -51,24 +50,22 @@ def find_response(response_message: RoborockMessage) -> None:
         if response is not None:
             future.set_result(response)
 
-    unsub = await mqtt_channel.subscribe(find_response)
+    unsub = await mqtt_channel.subscribe(on_message)
     try:
         await mqtt_channel.publish(roborock_message)
         return await asyncio.wait_for(future, timeout=_TIMEOUT)
-    except TimeoutError as ex:
-        raise RoborockException(timeout_error) from ex
     finally:
         unsub()
 
 
 async def send_decoded_command(
     mqtt_channel: MqttChannel,
     request_message: Q7RequestMessage,
-) -> dict[str, Any] | None:
+) -> dict[str, Any]:
     """Send a command on the MQTT channel and get a decoded response."""
     _LOGGER.debug("Sending B01 MQTT command: %s", request_message)
 
-    def find_response(response_message: RoborockMessage) -> Any | None:
+    def find_response(response_message: RoborockMessage) -> dict[str, Any] | None:
         """Handle incoming messages and resolve the future."""
         try:
             decoded_dps = decode_rpc_response(response_message)
@@ -94,15 +91,13 @@ def find_response(response_message: RoborockMessage) -> Any | None:
                 continue
             if isinstance(inner, dict) and inner.get("msgId") == str(request_message.msg_id):
                 _LOGGER.debug("Received query response: %s", inner)
-                # Check for error code (0 = success, non-zero = error)
                 code = inner.get("code", 0)
                 if code != 0:
                     error_msg = f"B01 command failed with code {code} ({request_message})"
                     _LOGGER.debug("B01 error response: %s", error_msg)
                     raise RoborockException(error_msg)
                 data = inner.get("data")
-                # All get commands should be dicts
-                if request_message.command.endswith(".get") and not isinstance(data, dict):
+                if not isinstance(data, dict):
                     raise RoborockException(f"Unexpected data type for response {data} ({request_message})")
                 return data
         return None
@@ -112,16 +107,16 @@ def find_response(response_message: RoborockMessage) -> Any | None:
             mqtt_channel,
             request_message,
             response_matcher=find_response,
-            timeout_error=f"B01 command timed out after {_TIMEOUT}s ({request_message})",
         )
+    except TimeoutError as ex:
+        raise RoborockException(f"B01 command timed out after {_TIMEOUT}s ({request_message})") from ex
     except RoborockException as ex:
         _LOGGER.warning(
             "Error sending B01 decoded command (%ss): %s",
             request_message,
             ex,
         )
         raise
-
     except Exception as ex:
         _LOGGER.exception(
             "Error sending B01 decoded command (%ss): %s",
@@ -138,9 +133,11 @@ async def send_map_command(mqtt_channel: MqttChannel, request_message: Q7Request
     raw ``MAP_RESPONSE`` payload bytes instead of a decoded RPC ``data`` payload.
     """
 
-    return await _send_command(
-        mqtt_channel,
-        request_message,
-        response_matcher=lambda response_message: _matches_map_response(response_message, version=B01_VERSION),
-        timeout_error=f"B01 map command timed out after {_TIMEOUT}s ({request_message})",
-    )
+    try:
+        return await _send_command(
+            mqtt_channel,
+            request_message,
+            response_matcher=lambda response_message: _matches_map_response(response_message, version=B01_VERSION),
+        )
+    except TimeoutError as ex:
+        raise RoborockException(f"B01 map command timed out after {_TIMEOUT}s ({request_message})") from ex

roborock/devices/traits/b01/q7/map.py

@@ -26,27 +26,28 @@ class Q7MapList(RoborockBase):
 
     map_list: list[Q7MapListEntry] = field(default_factory=list)
 
+    @property
+    def current_map_id(self) -> int | None:
+        """Current map id, preferring the entry marked current."""
+        if not self.map_list:
+            return None
+
+        ordered = sorted(self.map_list, key=lambda entry: entry.cur or False, reverse=True)
+        first = next(iter(ordered), None)
+        if first is None or not isinstance(first.id, int):
+            return None
+        return first.id
+
 
-class MapTrait(Trait):
+class MapTrait(Q7MapList, Trait):
     """Map retrieval + map metadata helpers for Q7 devices."""
 
     def __init__(self, channel: MqttChannel) -> None:
+        super().__init__()
         self._channel = channel
         # Map uploads are serialized per-device to avoid response cross-wiring.
         self._map_command_lock = asyncio.Lock()
-        self._map_list: Q7MapList | None = None
-
-    @property
-    def map_list(self) -> Q7MapList | None:
-        """Latest cached map list metadata, populated by ``refresh()``."""
-        return self._map_list
-
-    @property
-    def current_map_id(self) -> int | None:
-        """Current map id derived from cached map list metadata."""
-        if self._map_list is None:
-            return None
-        return self._extract_current_map_id(self._map_list)
+        self._loaded = False
 
     async def refresh(self) -> None:
         """Refresh cached map list metadata from the device."""
@@ -55,13 +56,15 @@ async def refresh(self) -> None:
             Q7RequestMessage(dps=B01_Q7_DPS, command=RoborockB01Q7Methods.GET_MAP_LIST, params={}),
         )
         if not isinstance(response, dict):
-            raise TypeError(f"Unexpected response type for GET_MAP_LIST: {type(response).__name__}: {response!r}")
+            raise RoborockException(
+                f"Unexpected response type for GET_MAP_LIST: {type(response).__name__}: {response!r}"
+            )
 
-        parsed = Q7MapList.from_dict(response)
-        if parsed is None:
-            raise TypeError(f"Failed to decode map list response: {response!r}")
+        if (parsed := Q7MapList.from_dict(response)) is None:
+            raise RoborockException(f"Failed to decode map list response: {response!r}")
 
-        self._map_list = parsed
+        self.map_list = parsed.map_list
+        self._loaded = True
 
     async def _get_map_payload(self, *, map_id: int) -> bytes:
         """Fetch raw map payload bytes for the given map id."""
@@ -75,25 +78,10 @@ async def _get_map_payload(self, *, map_id: int) -> bytes:
 
     async def get_current_map_payload(self) -> bytes:
         """Fetch raw map payload bytes for the currently selected map."""
-        if self._map_list is None:
+        if not self._loaded:
             await self.refresh()
 
         map_id = self.current_map_id
         if map_id is None:
-            raise RoborockException(f"Unable to determine map_id from map list response: {self._map_list!r}")
+            raise RoborockException(f"Unable to determine map_id from map list response: {self!r}")
         return await self._get_map_payload(map_id=map_id)
-
-    @staticmethod
-    def _extract_current_map_id(map_list_response: Q7MapList) -> int | None:
-        map_list = map_list_response.map_list
-        if not map_list:
-            return None
-
-        for entry in map_list:
-            if entry.cur and isinstance(entry.id, int):
-                return entry.id
-
-        first = map_list[0]
-        if isinstance(first.id, int):
-            return first.id
-        return None

tests/devices/traits/b01/q7/test_init.py

@@ -17,7 +17,7 @@
 from roborock.devices.traits.b01.q7 import Q7PropertiesApi
 from roborock.exceptions import RoborockException
 from roborock.protocols.b01_q7_protocol import B01_VERSION, Q7RequestMessage
-from roborock.roborock_message import RoborockB01Props, RoborockMessage, RoborockMessageProtocol
+from roborock.roborock_message import RoborockB01Props, RoborockMessageProtocol
 from tests.fixtures.channel_fixtures import FakeChannel
 
 from . import B01MessageBuilder
@@ -27,16 +27,12 @@ async def test_q7_api_query_values(
     q7_api: Q7PropertiesApi, fake_channel: FakeChannel, message_builder: B01MessageBuilder
 ):
     """Test that Q7PropertiesApi correctly converts raw values."""
-    # We need to construct the expected result based on the mappings
-    # status: 1 -> WAITING_FOR_ORDERS
-    # wind: 2 -> STANDARD
     response_data = {
         "status": 1,
         "wind": 2,
         "battery": 100,
     }
 
-    # Queue the response
     fake_channel.response_queue.append(message_builder.build(response_data))
 
     result = await q7_api.query_values(
@@ -48,23 +44,15 @@ async def test_q7_api_query_values(
 
     assert result is not None
     assert result.status == WorkStatusMapping.WAITING_FOR_ORDERS
-    # wind might be mapped to SCWindMapping.STANDARD (2)
-    # let's verify checking the prop definition in B01Props
-    # wind: SCWindMapping | None = None
-    # SCWindMapping.STANDARD is 2 ('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
     assert message.payload is not None
     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"]
@@ -110,8 +98,6 @@ async def test_send_decoded_command_non_dict_response(fake_channel: FakeChannel,
     message = message_builder.build("some_string_error")
     fake_channel.response_queue.append(message)
 
-    # Use a random string for command type to avoid needing import
-
     with pytest.raises(RoborockException, match="Unexpected data type for response"):
         await send_decoded_command(fake_channel, Q7RequestMessage(dps=10000, command="prop.get", params=[]))  # type: ignore[arg-type]
 
@@ -275,90 +261,3 @@ async def test_q7_api_clean_segments(
         "ctrl_value": SCDeviceCleanParam.START.code,
         "room_ids": [10, 11],
     }
-
-
-async def test_q7_api_get_current_map_payload(
-    q7_api: Q7PropertiesApi,
-    fake_channel: FakeChannel,
-    message_builder: B01MessageBuilder,
-):
-    """Fetch current map by map-list lookup, then upload_by_mapid."""
-    fake_channel.response_queue.append(message_builder.build({"map_list": [{"id": 1772093512, "cur": True}]}))
-    fake_channel.response_queue.append(
-        RoborockMessage(
-            protocol=RoborockMessageProtocol.MAP_RESPONSE,
-            payload=b"raw-map-payload",
-            version=b"B01",
-            seq=message_builder.seq + 1,
-        )
-    )
-
-    raw_payload = await q7_api.map.get_current_map_payload()
-    assert raw_payload == b"raw-map-payload"
-
-    assert len(fake_channel.published_messages) == 2
-
-    first = fake_channel.published_messages[0]
-    first_payload = json.loads(unpad(first.payload, AES.block_size))
-    assert first_payload["dps"]["10000"]["method"] == "service.get_map_list"
-    assert first_payload["dps"]["10000"]["params"] == {}
-
-    second = fake_channel.published_messages[1]
-    second_payload = json.loads(unpad(second.payload, AES.block_size))
-    assert second_payload["dps"]["10000"]["method"] == "service.upload_by_mapid"
-    assert second_payload["dps"]["10000"]["params"] == {"map_id": 1772093512}
-
-
-async def test_q7_api_map_trait_refresh_populates_cached_values(
-    q7_api: Q7PropertiesApi,
-    fake_channel: FakeChannel,
-    message_builder: B01MessageBuilder,
-):
-    """Map trait follows refresh + cached-value access pattern."""
-    fake_channel.response_queue.append(message_builder.build({"map_list": [{"id": 101, "cur": True}]}))
-
-    assert q7_api.map.map_list is None
-    assert q7_api.map.current_map_id is None
-
-    await q7_api.map.refresh()
-
-    assert len(fake_channel.published_messages) == 1
-    assert q7_api.map.map_list is not None
-    assert q7_api.map.map_list.map_list[0].id == 101
-    assert q7_api.map.map_list.map_list[0].cur is True
-    assert q7_api.map.current_map_id == 101
-
-
-async def test_q7_api_get_current_map_payload_falls_back_to_first_map(
-    q7_api: Q7PropertiesApi,
-    fake_channel: FakeChannel,
-    message_builder: B01MessageBuilder,
-):
-    """If no current map marker exists, first map in list is used."""
-    fake_channel.response_queue.append(message_builder.build({"map_list": [{"id": 111}, {"id": 222, "cur": False}]}))
-    fake_channel.response_queue.append(
-        RoborockMessage(
-            protocol=RoborockMessageProtocol.MAP_RESPONSE,
-            payload=b"raw-map-payload",
-            version=b"B01",
-            seq=message_builder.seq + 1,
-        )
-    )
-
-    await q7_api.map.get_current_map_payload()
-
-    second = fake_channel.published_messages[1]
-    second_payload = json.loads(unpad(second.payload, AES.block_size))
-    assert second_payload["dps"]["10000"]["params"] == {"map_id": 111}
-
-
-async def test_q7_api_get_current_map_payload_errors_without_map_list(
-    q7_api: Q7PropertiesApi,
-    fake_channel: FakeChannel,
-    message_builder: B01MessageBuilder,
-):
-    """Current-map payload fetch should fail clearly when map list is unusable."""
-    fake_channel.response_queue.append(message_builder.build({"map_list": []}))
-
-    with pytest.raises(RoborockException, match="Unable to determine map_id"):
-        await q7_api.map.get_current_map_payload()