q7/b01: command-layer segment clean + map payload retrieval helpers (split 1/3)

roborock/devices/rpc/b01_q7_channel.py

@@ -14,7 +14,7 @@
     decode_rpc_response,
     encode_mqtt_payload,
 )
-from roborock.roborock_message import RoborockMessage
+from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
 
 _LOGGER = logging.getLogger(__name__)
 _TIMEOUT = 10.0
@@ -99,3 +99,83 @@ def find_response(response_message: RoborockMessage) -> None:
         raise
     finally:
         unsub()
+
+
+async def send_map_command(mqtt_channel: MqttChannel, request_message: Q7RequestMessage) -> bytes:
+    """Send map upload command and wait for a map payload as bytes.
+
+    In practice, B01 map responses may arrive either as:
+    - a dedicated ``MAP_RESPONSE`` message with raw payload bytes, or
+    - a regular decoded RPC response that wraps a hex payload in ``data.payload``.
+
+    This helper accepts both response styles and returns the raw map payload bytes.
+    """
+
+    roborock_message = encode_mqtt_payload(request_message)
+    future: asyncio.Future[bytes] = asyncio.get_running_loop().create_future()
+    publish_started = asyncio.Event()
+
+    def find_response(response_message: RoborockMessage) -> None:
+        if future.done():
+            return
+
+        # Avoid accepting queued/stale MAP_RESPONSE messages before we actually
+        # publish this request.
+        if not publish_started.is_set():
+            return
+
+        if (
+            response_message.protocol == RoborockMessageProtocol.MAP_RESPONSE
+            and response_message.payload
+            and response_message.version == roborock_message.version
+        ):
+            future.set_result(response_message.payload)
+            return
+
+        try:
+            decoded_dps = decode_rpc_response(response_message)
+        except RoborockException:
+            return
+
+        for dps_value in decoded_dps.values():
+            if not isinstance(dps_value, str):
+                continue
+            try:
+                inner = json.loads(dps_value)
+            except (json.JSONDecodeError, TypeError):
+                continue
+            if not isinstance(inner, dict) or inner.get("msgId") != str(request_message.msg_id):
+                continue
+            code = inner.get("code", 0)
+            if code != 0:
+                if not future.done():
+                    future.set_exception(RoborockException(f"B01 command failed with code {code} ({request_message})"))
+                return
+            data = inner.get("data")
+            if isinstance(data, dict) and isinstance(data.get("payload"), str):
+                try:
+                    future.set_result(bytes.fromhex(data["payload"]))
+                except ValueError as ex:
+                    future.set_exception(
+                        RoborockException(
+                            f"Invalid hex payload in B01 map response: {data.get('payload')} ({request_message})"
+                        )
+                    )
+                    _LOGGER.debug(
+                        "Invalid hex payload in B01 map response (msgId=%s): %s (%s): %s",
+                        inner.get("msgId"),
+                        data.get("payload"),
+                        request_message,
+                        ex,
+                    )
+                return
+
+    unsub = await mqtt_channel.subscribe(find_response)
+    try:
+        publish_started.set()
+        await mqtt_channel.publish(roborock_message)
+        return await asyncio.wait_for(future, timeout=_TIMEOUT)
+    except TimeoutError as ex:
+        raise RoborockException(f"B01 map command timed out after {_TIMEOUT}s ({request_message})") from ex
+    finally:
+        unsub()

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

@@ -1,6 +1,7 @@
 """Traits for Q7 B01 devices.
 Potentially other devices may fall into this category in the future."""
 
+import asyncio
 from typing import Any
 
 from roborock import B01Props
@@ -13,9 +14,10 @@
     SCWindMapping,
     WaterLevelMapping,
 )
-from roborock.devices.rpc.b01_q7_channel import send_decoded_command
+from roborock.devices.rpc.b01_q7_channel import send_decoded_command, send_map_command
 from roborock.devices.traits import Trait
 from roborock.devices.transport.mqtt_channel import MqttChannel
+from roborock.exceptions import RoborockException
 from roborock.protocols.b01_q7_protocol import CommandType, ParamsType, Q7RequestMessage
 from roborock.roborock_message import RoborockB01Props
 from roborock.roborock_typing import RoborockB01Q7Methods
@@ -27,6 +29,8 @@
     "CleanSummaryTrait",
 ]
 
+_Q7_DPS = 10000
+
 
 class Q7PropertiesApi(Trait):
     """API for interacting with B01 devices."""
@@ -38,6 +42,8 @@ def __init__(self, channel: MqttChannel) -> None:
         """Initialize the B01Props API."""
         self._channel = channel
         self.clean_summary = CleanSummaryTrait(channel)
+        # Map uploads are serialized per-device to avoid response cross-wiring.
+        self._map_command_lock = asyncio.Lock()
 
     async def query_values(self, props: list[RoborockB01Props]) -> B01Props | None:
         """Query the device for the values of the given Q7 properties."""
@@ -87,6 +93,17 @@ async def start_clean(self) -> None:
             },
         )
 
+    async def clean_segments(self, segment_ids: list[int]) -> None:
+        """Start segment cleaning for the given ids (Q7 uses room ids)."""
+        await self.send(
+            command=RoborockB01Q7Methods.SET_ROOM_CLEAN,
+            params={
+                "clean_type": CleanTaskTypeMapping.ROOM.code,
+                "ctrl_value": SCDeviceCleanParam.START.code,
+                "room_ids": segment_ids,
+            },
+        )
+
     async def pause_clean(self) -> None:
         """Pause cleaning."""
         await self.send(
@@ -123,11 +140,61 @@ async def find_me(self) -> None:
             params={},
         )
 
+    async def get_map_list(self) -> dict[str, Any] | None:
+        """Return map list metadata from the robot."""
+        response = await self.send(
+            command=RoborockB01Q7Methods.GET_MAP_LIST,
+            params={},
+        )
+        if response is None:
+            return None
+        if not isinstance(response, dict):
+            raise TypeError(f"Unexpected response type for GET_MAP_LIST: {type(response).__name__}: {response!r}")
+        return response
+
+    async def get_current_map_id(self) -> int:
+        """Resolve and return the currently active map id."""
+        map_list_response = await self.get_map_list()
+        map_id = self._extract_current_map_id(map_list_response)
+        if map_id is None:
+            raise RoborockException(f"Unable to determine map_id from map list response: {map_list_response!r}")
+        return map_id
+
+    async def get_map_payload(self, *, map_id: int) -> bytes:
+        """Fetch raw map payload bytes for the given map id."""
+        request = Q7RequestMessage(
+            dps=_Q7_DPS,
+            command=RoborockB01Q7Methods.UPLOAD_BY_MAPID,
+            params={"map_id": map_id},
+        )
+        async with self._map_command_lock:
+            return await send_map_command(self._channel, request)
+
+    async def get_current_map_payload(self) -> bytes:
+        """Fetch raw map payload bytes for the map currently selected by the robot."""
+        return await self.get_map_payload(map_id=await self.get_current_map_id())
+
+    def _extract_current_map_id(self, map_list_response: dict[str, Any] | None) -> int | None:
+        if not isinstance(map_list_response, dict):
+            return None
+        map_list = map_list_response.get("map_list")
+        if not isinstance(map_list, list) or not map_list:
+            return None
+
+        for entry in map_list:
+            if isinstance(entry, dict) and entry.get("cur") and isinstance(entry.get("id"), int):
+                return entry["id"]
+
+        first = map_list[0]
+        if isinstance(first, dict) and isinstance(first.get("id"), int):
+            return first["id"]
+        return None
+
     async def send(self, command: CommandType, params: ParamsType) -> Any:
         """Send a command to the device."""
         return await send_decoded_command(
             self._channel,
-            Q7RequestMessage(dps=10000, command=command, params=params),
+            Q7RequestMessage(dps=_Q7_DPS, command=command, params=params),
         )
 
 

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, RoborockMessageProtocol
+from roborock.roborock_message import RoborockB01Props, RoborockMessage, RoborockMessageProtocol
 from tests.fixtures.channel_fixtures import FakeChannel
 
 from . import B01MessageBuilder
@@ -257,3 +257,114 @@ async def test_q7_api_find_me(q7_api: Q7PropertiesApi, fake_channel: FakeChannel
     payload_data = json.loads(unpad(message.payload, AES.block_size))
     assert payload_data["dps"]["10000"]["method"] == "service.find_device"
     assert payload_data["dps"]["10000"]["params"] == {}
+
+
+async def test_q7_api_clean_segments(
+    q7_api: Q7PropertiesApi, fake_channel: FakeChannel, message_builder: B01MessageBuilder
+):
+    """Test room/segment cleaning helper for Q7."""
+    fake_channel.response_queue.append(message_builder.build({"result": "ok"}))
+    await q7_api.clean_segments([10, 11])
+
+    assert len(fake_channel.published_messages) == 1
+    message = fake_channel.published_messages[0]
+    payload_data = json.loads(unpad(message.payload, AES.block_size))
+    assert payload_data["dps"]["10000"]["method"] == "service.set_room_clean"
+    assert payload_data["dps"]["10000"]["params"] == {
+        "clean_type": CleanTaskTypeMapping.ROOM.code,
+        "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.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_get_current_map_payload_rpc_wrapped_hex_payload(
+    q7_api: Q7PropertiesApi,
+    fake_channel: FakeChannel,
+    message_builder: B01MessageBuilder,
+):
+    """Fetch current map via RPC-wrapped hex payload response."""
+    fake_channel.response_queue.append(message_builder.build({"map_list": [{"id": 1772093512, "cur": True}]}))
+    fake_channel.response_queue.append(message_builder.build({"payload": "68656c6c6f"}))
+
+    raw_payload = await q7_api.get_current_map_payload()
+    assert raw_payload == b"hello"
+
+
+async def test_q7_api_get_current_map_payload_rpc_wrapped_invalid_hex_errors(
+    q7_api: Q7PropertiesApi,
+    fake_channel: FakeChannel,
+    message_builder: B01MessageBuilder,
+):
+    """Invalid hex payload should fail fast (not time out)."""
+    fake_channel.response_queue.append(message_builder.build({"map_list": [{"id": 1772093512, "cur": True}]}))
+    fake_channel.response_queue.append(message_builder.build({"payload": "not-hex"}))
+
+    with pytest.raises(RoborockException, match="Invalid hex payload"):
+        await q7_api.get_current_map_payload()
+
+
+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.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.get_current_map_payload()