From d9211b2558e1110b0d9c5611b27591fe2b241584 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Wed, 1 Apr 2026 17:24:42 +0200 Subject: [PATCH 1/4] feat: implement map-related commands and payload decoding for B01/Q7 devices --- roborock/devices/rpc/b01_q7_channel.py | 46 +++++-- roborock/devices/traits/b01/q7/__init__.py | 15 +- roborock/devices/traits/b01/q7/map.py | 34 +---- roborock/devices/traits/b01/q7/map_content.py | 56 ++++---- roborock/map/b01_map_parser.py | 70 ++-------- roborock/protocols/b01_q7_protocol.py | 50 +++++++ tests/devices/traits/b01/q7/__init__.py | 10 ++ tests/devices/traits/b01/q7/conftest.py | 4 +- tests/devices/traits/b01/q7/test_map.py | 75 ---------- .../devices/traits/b01/q7/test_map_content.py | 128 +++++++++++------- tests/map/test_b01_map_parser.py | 12 +- 11 files changed, 232 insertions(+), 268 deletions(-) diff --git a/roborock/devices/rpc/b01_q7_channel.py b/roborock/devices/rpc/b01_q7_channel.py index d2e373d3..44dc60e6 100644 --- a/roborock/devices/rpc/b01_q7_channel.py +++ b/roborock/devices/rpc/b01_q7_channel.py @@ -10,7 +10,14 @@ from roborock.devices.transport.mqtt_channel import MqttChannel from roborock.exceptions import RoborockException -from roborock.protocols.b01_q7_protocol import B01_VERSION, Q7RequestMessage, decode_rpc_response, encode_mqtt_payload +from roborock.protocols.b01_q7_protocol import ( + B01_VERSION, + MapKey, + Q7RequestMessage, + decode_map_payload, + decode_rpc_response, + encode_mqtt_payload, +) from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol _LOGGER = logging.getLogger(__name__) @@ -127,18 +134,29 @@ def find_response(response_message: RoborockMessage) -> DecodedB01Response | Non raise -async def send_map_command(mqtt_channel: MqttChannel, request_message: Q7RequestMessage) -> bytes: - """Send map upload command and wait for MAP_RESPONSE payload bytes. +class MapRpcChannel: + """RPC channel for map-related commands on B01/Q7 devices.""" - This stays separate from ``send_decoded_command()`` because map uploads arrive as - raw ``MAP_RESPONSE`` payload bytes instead of a decoded RPC ``data`` payload. - """ + def __init__(self, mqtt_channel: MqttChannel, map_key: MapKey) -> None: + self._mqtt_channel = mqtt_channel + self._map_key = map_key - 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 + async def send_map_command(self, request_message: Q7RequestMessage) -> bytes: + """Send map upload command and wait for MAP_RESPONSE payload bytes. + + This stays separate from ``send_decoded_command()`` because map uploads arrive as + raw ``MAP_RESPONSE`` payload bytes instead of a decoded RPC ``data`` payload. + + The response is a protocol buffer than can be parsed by the map parser library. + """ + + try: + raw_payload = await _send_command( + self._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 + + return decode_map_payload(raw_payload, map_key=self._map_key) diff --git a/roborock/devices/traits/b01/q7/__init__.py b/roborock/devices/traits/b01/q7/__init__.py index f29f287b..e00d540b 100644 --- a/roborock/devices/traits/b01/q7/__init__.py +++ b/roborock/devices/traits/b01/q7/__init__.py @@ -18,10 +18,10 @@ SCWindMapping, WaterLevelMapping, ) -from roborock.devices.rpc.b01_q7_channel import send_decoded_command +from roborock.devices.rpc.b01_q7_channel import MapRpcChannel, send_decoded_command from roborock.devices.traits import Trait from roborock.devices.transport.mqtt_channel import MqttChannel -from roborock.protocols.b01_q7_protocol import B01_Q7_DPS, CommandType, ParamsType, Q7RequestMessage +from roborock.protocols.b01_q7_protocol import B01_Q7_DPS, CommandType, ParamsType, Q7RequestMessage, create_map_key from roborock.roborock_message import RoborockB01Props from roborock.roborock_typing import RoborockB01Q7Methods @@ -51,9 +51,12 @@ class Q7PropertiesApi(Trait): map_content: MapContentTrait """Trait for fetching parsed current map content.""" - def __init__(self, channel: MqttChannel, *, device: HomeDataDevice, product: HomeDataProduct) -> None: + def __init__( + self, channel: MqttChannel, map_rpc_channel: MapRpcChannel, device: HomeDataDevice, product: HomeDataProduct + ) -> None: """Initialize the Q7 API.""" self._channel = channel + self._map_rpc_channel = map_rpc_channel self._device = device self._product = product @@ -63,9 +66,8 @@ def __init__(self, channel: MqttChannel, *, device: HomeDataDevice, product: Hom self.clean_summary = CleanSummaryTrait(channel) self.map = MapTrait(channel) self.map_content = MapContentTrait( + self._map_rpc_channel, self.map, - serial=device.sn, - model=product.model, ) async def query_values(self, props: list[RoborockB01Props]) -> B01Props | None: @@ -173,4 +175,5 @@ async def send(self, command: CommandType, params: ParamsType) -> Any: def create(product: HomeDataProduct, device: HomeDataDevice, channel: MqttChannel) -> Q7PropertiesApi: """Create traits for B01 Q7 devices.""" - return Q7PropertiesApi(channel, device=device, product=product) + map_rpc_channel = MapRpcChannel(channel, map_key=create_map_key(serial=device.sn or "", model=product.model or "")) + return Q7PropertiesApi(channel, device=device, product=product, map_rpc_channel=map_rpc_channel) diff --git a/roborock/devices/traits/b01/q7/map.py b/roborock/devices/traits/b01/q7/map.py index c48379c2..f367e407 100644 --- a/roborock/devices/traits/b01/q7/map.py +++ b/roborock/devices/traits/b01/q7/map.py @@ -1,9 +1,7 @@ """Map trait for B01 Q7 devices.""" -import asyncio - from roborock.data import Q7MapList -from roborock.devices.rpc.b01_q7_channel import send_decoded_command, send_map_command +from roborock.devices.rpc.b01_q7_channel import send_decoded_command from roborock.devices.traits import Trait from roborock.devices.transport.mqtt_channel import MqttChannel from roborock.exceptions import RoborockException @@ -12,14 +10,15 @@ class MapTrait(Q7MapList, Trait): - """Map retrieval + map metadata helpers for Q7 devices.""" + """Map trait for B01/Q7 devices, responsible for fetching and caching map list metadata. + + The MapContent is fetched from the MapContent trait, which relies on this trait to determine the + current map ID to fetch. + """ 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._loaded = False async def refresh(self) -> None: """Refresh cached map list metadata from the device.""" @@ -36,24 +35,3 @@ async def refresh(self) -> None: raise RoborockException(f"Failed to decode map list response: {response!r}") 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.""" - request = Q7RequestMessage( - dps=B01_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 currently selected map.""" - 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!r}") - return await self._get_map_payload(map_id=map_id) diff --git a/roborock/devices/traits/b01/q7/map_content.py b/roborock/devices/traits/b01/q7/map_content.py index db00119b..40e4f781 100644 --- a/roborock/devices/traits/b01/q7/map_content.py +++ b/roborock/devices/traits/b01/q7/map_content.py @@ -8,14 +8,18 @@ For B01/Q7 devices, the underlying raw map payload is retrieved via `MapTrait`. """ +import asyncio from dataclasses import dataclass from vacuum_map_parser_base.map_data import MapData from roborock.data import RoborockBase +from roborock.devices.rpc.b01_q7_channel import MapRpcChannel from roborock.devices.traits import Trait from roborock.exceptions import RoborockException from roborock.map.b01_map_parser import B01MapParser, B01MapParserConfig +from roborock.protocols.b01_q7_protocol import B01_Q7_DPS, Q7RequestMessage +from roborock.roborock_typing import RoborockB01Q7Methods from .map import MapTrait @@ -51,48 +55,46 @@ class MapContentTrait(MapContent, Trait): def __init__( self, + map_rpc_channel: MapRpcChannel, map_trait: MapTrait, *, - serial: str, - model: str, map_parser_config: B01MapParserConfig | None = None, ) -> None: super().__init__() + self._map_rpc_channel = map_rpc_channel self._map_trait = map_trait - self._serial = serial - self._model = model self._map_parser = B01MapParser(map_parser_config) + # Map uploads are serialized per-device to avoid response cross-wiring. + self._map_command_lock = asyncio.Lock() async def refresh(self) -> None: - """Fetch, decode, and parse the current map payload.""" - raw_payload = await self._map_trait.get_current_map_payload() - parsed = self.parse_map_content(raw_payload) - self.image_content = parsed.image_content - self.map_data = parsed.map_data - self.raw_api_response = parsed.raw_api_response - - def parse_map_content(self, response: bytes) -> MapContent: - """Parse map content from raw bytes. - - This mirrors the v1 trait behavior so cached map payload bytes can be - reparsed without going back to the device. + """Fetch, decode, and parse the current map payload. + + This relies on the Map Trait already having fetched the map list metadata + so it can determine the current map_id. """ + # Users must call first + if (map_id := self._map_trait.current_map_id) is None: + raise RoborockException("Unable to determine current map ID") + + request = Q7RequestMessage( + dps=B01_Q7_DPS, + command=RoborockB01Q7Methods.UPLOAD_BY_MAPID, + params={"map_id": map_id}, + ) + async with self._map_command_lock: + raw_payload = await self._map_rpc_channel.send_map_command(request) + try: - parsed_data = self._map_parser.parse( - response, - serial=self._serial, - model=self._model, - ) + parsed_data = self._map_parser.parse(raw_payload) except RoborockException: raise except Exception as ex: - raise RoborockException("Failed to parse B01 map data") from ex + raise RoborockException(f"Uncaught exception parsing B01 map data: {ex}") from ex if parsed_data.image_content is None: raise RoborockException("Failed to render B01 map image") - return MapContent( - image_content=parsed_data.image_content, - map_data=parsed_data.map_data, - raw_api_response=response, - ) + self.image_content = parsed_data.image_content + self.map_data = parsed_data.map_data + self.raw_api_response = raw_payload diff --git a/roborock/map/b01_map_parser.py b/roborock/map/b01_map_parser.py index 5e0a4bad..b57912e4 100644 --- a/roborock/map/b01_map_parser.py +++ b/roborock/map/b01_map_parser.py @@ -1,31 +1,19 @@ """Module for parsing B01/Q7 map content. -Observed Q7 `MAP_RESPONSE` payloads follow this decode pipeline: -- base64-encoded ASCII -- AES-ECB encrypted with the derived map key -- PKCS7 padded -- ASCII hex for a zlib-compressed SCMap payload - The inner SCMap blob is parsed with protobuf messages generated from `roborock/map/proto/b01_scmap.proto`. """ -import base64 -import binascii -import hashlib import io -import zlib from dataclasses import dataclass -from Crypto.Cipher import AES -from google.protobuf.message import DecodeError, Message +from google.protobuf.message import DecodeError from PIL import Image from vacuum_map_parser_base.config.image_config import ImageConfig from vacuum_map_parser_base.map_data import ImageData, MapData from roborock.exceptions import RoborockException from roborock.map.proto.b01_scmap_pb2 import RobotMap # type: ignore[attr-defined] -from roborock.protocol import Utils from .map_parser import ParsedMapData @@ -46,10 +34,9 @@ class B01MapParser: def __init__(self, config: B01MapParserConfig | None = None) -> None: self._config = config or B01MapParserConfig() - def parse(self, raw_payload: bytes, *, serial: str, model: str) -> ParsedMapData: - """Parse a raw MAP_RESPONSE payload and return a PNG + MapData.""" - inflated = _decode_b01_map_payload(raw_payload, serial=serial, model=model) - parsed = _parse_scmap_payload(inflated) + def parse(self, payload: bytes) -> ParsedMapData: + """Parse an inflated SCMap payload and return a PNG + MapData.""" + parsed = _parse_scmap_payload(payload) size_x, size_y, grid = _extract_grid(parsed) room_names = _extract_room_names(parsed) @@ -78,54 +65,13 @@ def parse(self, raw_payload: bytes, *, serial: str, model: str) -> ParsedMapData ) -def _derive_map_key(serial: str, model: str) -> bytes: - """Derive the B01/Q7 map decrypt key from serial + model.""" - model_suffix = model.split(".")[-1] - model_key = (model_suffix + "0" * 16)[:16].encode() - material = f"{serial}+{model_suffix}+{serial}".encode() - encrypted = Utils.encrypt_ecb(material, model_key) - md5 = hashlib.md5(base64.b64encode(encrypted), usedforsecurity=False).hexdigest() - return md5[8:24].encode() - - -def _decode_base64_payload(raw_payload: bytes) -> bytes: - blob = raw_payload.strip() - padded = blob + b"=" * (-len(blob) % 4) - try: - return base64.b64decode(padded, validate=True) - except binascii.Error as err: - raise RoborockException("Failed to decode B01 map payload") from err - - -def _decode_b01_map_payload(raw_payload: bytes, *, serial: str, model: str) -> bytes: - """Decode raw B01 `MAP_RESPONSE` payload into inflated SCMap bytes.""" - # TODO: Move this lower-level B01 transport decode under `roborock.protocols` - # so this module only handles SCMap parsing/rendering. - encrypted_payload = _decode_base64_payload(raw_payload) - if len(encrypted_payload) % AES.block_size != 0: - raise RoborockException("Unexpected encrypted B01 map payload length") - - map_key = _derive_map_key(serial, model) - - try: - compressed_hex = Utils.decrypt_ecb(encrypted_payload, map_key).decode("ascii") - compressed_payload = bytes.fromhex(compressed_hex) - return zlib.decompress(compressed_payload) - except (ValueError, UnicodeDecodeError, zlib.error) as err: - raise RoborockException("Failed to decode B01 map payload") from err - - -def _parse_proto(blob: bytes, message: Message, *, context: str) -> None: - try: - message.ParseFromString(blob) - except DecodeError as err: - raise RoborockException(f"Failed to parse {context}") from err - - def _parse_scmap_payload(payload: bytes) -> RobotMap: """Parse inflated SCMap bytes into a generated protobuf message.""" parsed = RobotMap() - _parse_proto(payload, parsed, context="B01 SCMap") + try: + parsed.ParseFromString(payload) + except DecodeError as err: + raise RoborockException("Failed to parse B01 SCMap") from err return parsed diff --git a/roborock/protocols/b01_q7_protocol.py b/roborock/protocols/b01_q7_protocol.py index 0d109d78..be4a0122 100644 --- a/roborock/protocols/b01_q7_protocol.py +++ b/roborock/protocols/b01_q7_protocol.py @@ -1,7 +1,11 @@ """Roborock B01 Protocol encoding and decoding.""" +import base64 +import binascii +import hashlib import json import logging +import zlib from dataclasses import dataclass, field from typing import Any @@ -10,6 +14,7 @@ from roborock import RoborockB01Q7Methods from roborock.exceptions import RoborockException +from roborock.protocol import Utils from roborock.roborock_message import ( RoborockMessage, RoborockMessageProtocol, @@ -80,3 +85,48 @@ def decode_rpc_response(message: RoborockMessage) -> dict[int, Any]: return {int(key): value for key, value in datapoints.items()} except ValueError: raise RoborockException(f"Invalid B01 message format: 'dps' key should be an integer for {message.payload!r}") + + +@dataclass +class MapKey: + """Data class for holding a B01 map decryption key.""" + + key: bytes + + +def create_map_key(serial: str, model: str) -> MapKey: + """Derive the B01/Q7 map decrypt key from serial + model.""" + model_suffix = model.split(".")[-1] + model_key = (model_suffix + "0" * 16)[:16].encode() + material = f"{serial}+{model_suffix}+{serial}".encode() + encrypted = Utils.encrypt_ecb(material, model_key) + md5 = hashlib.md5(base64.b64encode(encrypted), usedforsecurity=False).hexdigest() + return MapKey(key=md5[8:24].encode()) + + +def decode_map_payload(raw_payload: bytes, map_key: MapKey) -> bytes: + """Decode raw B01 `MAP_RESPONSE` payload into inflated SCMap bytes.""" + encrypted_payload = _decode_base64_payload(raw_payload) + payload_len = len(encrypted_payload) + if payload_len % AES.block_size != 0: + raise RoborockException( + f"Unexpected encrypted B01 map payload length: {payload_len} (not a multiple of AES block size)" + ) + + try: + compressed_hex = Utils.decrypt_ecb(encrypted_payload, token=map_key.key).decode("ascii") + compressed_payload = bytes.fromhex(compressed_hex) + return zlib.decompress(compressed_payload) + except (ValueError, UnicodeDecodeError, zlib.error) as err: + raise RoborockException("Failed to decode B01 map payload") from err + + +def _decode_base64_payload(raw_payload: bytes) -> bytes: + """Decode base64 payload.""" + + blob = raw_payload.strip() + padded = blob + b"=" * (-len(blob) % 4) + try: + return base64.b64decode(padded, validate=True) + except binascii.Error as err: + raise RoborockException("Failed to decode B01 map payload") from err diff --git a/tests/devices/traits/b01/q7/__init__.py b/tests/devices/traits/b01/q7/__init__.py index d79a65d7..128a0924 100644 --- a/tests/devices/traits/b01/q7/__init__.py +++ b/tests/devices/traits/b01/q7/__init__.py @@ -39,3 +39,13 @@ def _build_dps(self, message: dict[str, Any] | str) -> RoborockMessage: version=b"B01", seq=self.seq, ) + + def build_map_response(self, payload: bytes) -> RoborockMessage: + """Build a dummy MAP_RESPONSE message.""" + self.seq += 1 + return RoborockMessage( + protocol=RoborockMessageProtocol.MAP_RESPONSE, + payload=payload, + version=b"B01", + seq=self.seq, + ) diff --git a/tests/devices/traits/b01/q7/conftest.py b/tests/devices/traits/b01/q7/conftest.py index 4f55ba0f..caf2097d 100644 --- a/tests/devices/traits/b01/q7/conftest.py +++ b/tests/devices/traits/b01/q7/conftest.py @@ -6,7 +6,7 @@ import pytest from roborock.data import HomeDataDevice, HomeDataProduct, RoborockCategory -from roborock.devices.traits.b01.q7 import Q7PropertiesApi +from roborock.devices.traits.b01.q7 import Q7PropertiesApi, create from tests.fixtures.channel_fixtures import FakeChannel from . import B01MessageBuilder @@ -40,7 +40,7 @@ def device_fixture() -> HomeDataDevice: @pytest.fixture(name="q7_api") def q7_api_fixture(fake_channel: FakeChannel, device: HomeDataDevice, product: HomeDataProduct) -> Q7PropertiesApi: - return Q7PropertiesApi(fake_channel, device=device, product=product) # type: ignore[arg-type] + return create(product, device, fake_channel) # type: ignore[arg-type] @pytest.fixture(name="expected_msg_id", autouse=True) diff --git a/tests/devices/traits/b01/q7/test_map.py b/tests/devices/traits/b01/q7/test_map.py index d158d363..04b6a7df 100644 --- a/tests/devices/traits/b01/q7/test_map.py +++ b/tests/devices/traits/b01/q7/test_map.py @@ -1,50 +1,10 @@ -import json - -import pytest -from Crypto.Cipher import AES -from Crypto.Util.Padding import unpad - from roborock.data import Q7MapList, Q7MapListEntry from roborock.devices.traits.b01.q7 import Q7PropertiesApi -from roborock.exceptions import RoborockException -from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol from tests.fixtures.channel_fixtures import FakeChannel from . import B01MessageBuilder -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, @@ -64,41 +24,6 @@ async def test_q7_api_map_trait_refresh_populates_cached_values( 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() - - def test_q7_map_list_current_map_id_prefers_marked_current(): """Current-map resolution prefers the entry marked current.""" map_list = Q7MapList( diff --git a/tests/devices/traits/b01/q7/test_map_content.py b/tests/devices/traits/b01/q7/test_map_content.py index c0b62a4f..7ec11112 100644 --- a/tests/devices/traits/b01/q7/test_map_content.py +++ b/tests/devices/traits/b01/q7/test_map_content.py @@ -1,13 +1,14 @@ -from typing import cast +import json from unittest.mock import patch import pytest +from Crypto.Cipher import AES +from Crypto.Util.Padding import unpad from vacuum_map_parser_base.map_data import MapData from roborock.devices.traits.b01.q7 import Q7PropertiesApi -from roborock.devices.transport.mqtt_channel import MqttChannel from roborock.exceptions import RoborockException -from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol +from roborock.map.b01_map_parser import ParsedMapData from tests.fixtures.channel_fixtures import FakeChannel from . import B01MessageBuilder @@ -18,62 +19,91 @@ async def test_q7_map_content_refresh_populates_cached_values( fake_channel: FakeChannel, message_builder: B01MessageBuilder, ): - 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, - ) + fake_channel.response_queue.extend( + [ + message_builder.build({"map_list": [{"id": 1772093512, "cur": True}]}), + message_builder.build_map_response(b"raw-map-payload"), + ] ) + # Ensure we have map metadata first + await q7_api.map.refresh() + dummy_map_data = MapData() - with patch( - "roborock.devices.traits.b01.q7.map_content.B01MapParser.parse", - return_value=type("X", (), {"image_content": b"pngbytes", "map_data": dummy_map_data})(), - ) as parse: + parsed_map_data = ParsedMapData( + image_content=b"pngbytes", + map_data=dummy_map_data, + ) + with ( + patch( + "roborock.devices.rpc.b01_q7_channel.decode_map_payload", + return_value=b"inflated-payload", + ), + patch( + "roborock.devices.traits.b01.q7.map_content.B01MapParser.parse", + return_value=parsed_map_data, + ) as parse, + ): await q7_api.map_content.refresh() assert q7_api.map_content.image_content == b"pngbytes" assert q7_api.map_content.map_data is dummy_map_data - assert q7_api.map_content.raw_api_response == b"raw-map-payload" + assert q7_api.map_content.raw_api_response == b"inflated-payload" - parse.assert_called_once() + parse.assert_called_once_with(b"inflated-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" -def test_q7_map_content_parse_errors_cleanly(q7_api: Q7PropertiesApi): - with patch("roborock.devices.traits.b01.q7.map_content.B01MapParser.parse", side_effect=ValueError("boom")): - with pytest.raises(RoborockException, match="Failed to parse B01 map data"): - q7_api.map_content.parse_map_content(b"raw") + 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} -def test_q7_map_content_preserves_specific_roborock_errors(q7_api: Q7PropertiesApi): - with patch( - "roborock.devices.traits.b01.q7.map_content.B01MapParser.parse", - side_effect=RoborockException("Specific decoder failure"), +async def test_q7_map_content_refresh_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.extend( + [ + message_builder.build({"map_list": [{"id": 111}, {"id": 222, "cur": False}]}), + message_builder.build_map_response(b"raw-map-payload"), + ] + ) + + # Load current map + await q7_api.map.refresh() + + dummy_map_data = MapData() + with ( + patch( + "roborock.devices.rpc.b01_q7_channel.decode_map_payload", + return_value=b"inflated-payload", + ), + patch( + "roborock.devices.traits.b01.q7.map_content.B01MapParser.parse", + return_value=type("X", (), {"image_content": b"pngbytes", "map_data": dummy_map_data})(), + ), ): - with pytest.raises(RoborockException, match="Specific decoder failure"): - q7_api.map_content.parse_map_content(b"raw") - - -def test_q7_map_content_requires_metadata_at_init(fake_channel: FakeChannel): - from roborock.data import HomeDataDevice, HomeDataProduct, RoborockCategory - - with pytest.raises(ValueError, match="requires device serial number and product model metadata"): - Q7PropertiesApi( - cast(MqttChannel, fake_channel), - device=HomeDataDevice( - duid="abc123", - name="Q7", - local_key="key123key123key1", - product_id="product-id-q7", - sn=None, - ), - product=HomeDataProduct( - id="product-id-q7", - name="Roborock Q7", - model="roborock.vacuum.sc05", - category=RoborockCategory.VACUUM, - ), - ) + await q7_api.map_content.refresh() + + 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_map_content_refresh_errors_without_map_list( + q7_api: Q7PropertiesApi, + fake_channel: FakeChannel, + message_builder: B01MessageBuilder, +): + """Refresh 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 current map ID"): + await q7_api.map_content.refresh() diff --git a/tests/map/test_b01_map_parser.py b/tests/map/test_b01_map_parser.py index 873e3005..0829182e 100644 --- a/tests/map/test_b01_map_parser.py +++ b/tests/map/test_b01_map_parser.py @@ -13,6 +13,7 @@ from roborock.exceptions import RoborockException from roborock.map.b01_map_parser import B01MapParser, _parse_scmap_payload from roborock.map.proto.b01_scmap_pb2 import RobotMap # type: ignore[attr-defined] +from roborock.protocols.b01_q7_protocol import create_map_key, decode_map_payload FIXTURE = Path(__file__).resolve().parent / "testdata" / "raw-mqtt-map301.bin.inflated.bin.gz" @@ -32,12 +33,13 @@ def test_b01_map_parser_decodes_and_renders_fixture() -> None: inflated = gzip.decompress(FIXTURE.read_bytes()) compressed = zlib.compress(inflated) - map_key = _derive_map_key(serial, model) - encrypted = AES.new(map_key, AES.MODE_ECB).encrypt(pad(compressed.hex().encode(), AES.block_size)) + map_key = create_map_key(serial, model) + encrypted = AES.new(map_key.key, AES.MODE_ECB).encrypt(pad(compressed.hex().encode(), AES.block_size)) payload = base64.b64encode(encrypted) parser = B01MapParser() - parsed = parser.parse(payload, serial=serial, model=model) + inflated_payload = decode_map_payload(payload, map_key=map_key) + parsed = parser.parse(inflated_payload) assert parsed.image_content is not None assert parsed.image_content.startswith(b"\x89PNG\r\n\x1a\n") @@ -126,5 +128,5 @@ def test_b01_scmap_parser_maps_observed_schema_fields() -> None: def test_b01_map_parser_rejects_invalid_payload() -> None: parser = B01MapParser() - with pytest.raises(RoborockException, match="Failed to decode B01 map payload"): - parser.parse(b"not a map", serial="testsn012345", model="roborock.vacuum.sc05") + with pytest.raises(RoborockException, match="Failed to parse B01 SCMap"): + parser.parse(b"not a map") From ea5d444977c3a0943c096a310f39d166773c6584 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Fri, 3 Apr 2026 17:38:02 +0200 Subject: [PATCH 2/4] Update roborock/devices/traits/b01/q7/map_content.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- roborock/devices/traits/b01/q7/map_content.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/roborock/devices/traits/b01/q7/map_content.py b/roborock/devices/traits/b01/q7/map_content.py index 40e4f781..afc36d02 100644 --- a/roborock/devices/traits/b01/q7/map_content.py +++ b/roborock/devices/traits/b01/q7/map_content.py @@ -90,7 +90,7 @@ async def refresh(self) -> None: except RoborockException: raise except Exception as ex: - raise RoborockException(f"Uncaught exception parsing B01 map data: {ex}") from ex + raise RoborockException("Failed to parse B01 map data") from ex if parsed_data.image_content is None: raise RoborockException("Failed to render B01 map image") From 6ac8462c28c522e320119dd43b5176d234da0612 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Fri, 3 Apr 2026 17:38:22 +0200 Subject: [PATCH 3/4] feat: Update roborock/devices/rpc/b01_q7_channel.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- roborock/devices/rpc/b01_q7_channel.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/roborock/devices/rpc/b01_q7_channel.py b/roborock/devices/rpc/b01_q7_channel.py index 44dc60e6..b9926e3e 100644 --- a/roborock/devices/rpc/b01_q7_channel.py +++ b/roborock/devices/rpc/b01_q7_channel.py @@ -142,12 +142,15 @@ def __init__(self, mqtt_channel: MqttChannel, map_key: MapKey) -> None: self._map_key = map_key async def send_map_command(self, request_message: Q7RequestMessage) -> bytes: - """Send map upload command and wait for MAP_RESPONSE payload bytes. + """Send a map upload command and return decoded SCMap bytes. - This stays separate from ``send_decoded_command()`` because map uploads arrive as - raw ``MAP_RESPONSE`` payload bytes instead of a decoded RPC ``data`` payload. + This publishes the request and waits for a matching ``MAP_RESPONSE`` message + with the correct protocol version. The raw ``MAP_RESPONSE`` payload bytes are + then decoded/inflated via :func:`decode_map_payload` using this channel's + ``map_key``, and the resulting SCMap bytes are returned. - The response is a protocol buffer than can be parsed by the map parser library. + The returned value is the decoded map data bytes suitable for passing to the + map parser library, not the raw MQTT ``MAP_RESPONSE`` payload bytes. """ try: From 4bd97fe444c07403d9b3dceb40b04c5f60ad0240 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Fri, 3 Apr 2026 17:39:34 +0200 Subject: [PATCH 4/4] fix: ensure device serial number and product model are provided when creating Q7PropertiesApi --- roborock/devices/traits/b01/q7/__init__.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/roborock/devices/traits/b01/q7/__init__.py b/roborock/devices/traits/b01/q7/__init__.py index e00d540b..3a3707d8 100644 --- a/roborock/devices/traits/b01/q7/__init__.py +++ b/roborock/devices/traits/b01/q7/__init__.py @@ -21,6 +21,7 @@ from roborock.devices.rpc.b01_q7_channel import MapRpcChannel, send_decoded_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 B01_Q7_DPS, CommandType, ParamsType, Q7RequestMessage, create_map_key from roborock.roborock_message import RoborockB01Props from roborock.roborock_typing import RoborockB01Q7Methods @@ -175,5 +176,9 @@ async def send(self, command: CommandType, params: ParamsType) -> Any: def create(product: HomeDataProduct, device: HomeDataDevice, channel: MqttChannel) -> Q7PropertiesApi: """Create traits for B01 Q7 devices.""" - map_rpc_channel = MapRpcChannel(channel, map_key=create_map_key(serial=device.sn or "", model=product.model or "")) + if device.sn is None or product.model is None: + raise RoborockException( + f"Device serial number and product model are required (sn:: {device.sn}, model: {product.model})" + ) + map_rpc_channel = MapRpcChannel(channel, map_key=create_map_key(serial=device.sn, model=product.model)) return Q7PropertiesApi(channel, device=device, product=product, map_rpc_channel=map_rpc_channel)