From 7bfe6996fe6247c01bd61762eb9688e52a92a2e9 Mon Sep 17 00:00:00 2001 From: Harry Coureau Date: Sat, 7 Nov 2026 19:01:07 +0000 Subject: [PATCH] feat: add Q10 map management controls --- .../data/b01_q10/b01_q10_code_mappings.py | 11 +- roborock/devices/traits/b01/q10/map.py | 68 ++++++- roborock/devices/traits/b01/q10/maps.py | 24 ++- roborock/map/b01_q10_overlays.py | 30 ++- roborock/protocols/b01_q10_protocol.py | 118 ++++++++++- tests/devices/traits/b01/q10/test_map.py | 178 ++++++++++++++++- tests/protocols/test_b01_q10_protocol.py | 189 +++++++++++++++++- 7 files changed, 606 insertions(+), 12 deletions(-) diff --git a/roborock/data/b01_q10/b01_q10_code_mappings.py b/roborock/data/b01_q10/b01_q10_code_mappings.py index a35f81f24..afbbff451 100644 --- a/roborock/data/b01_q10/b01_q10_code_mappings.py +++ b/roborock/data/b01_q10/b01_q10_code_mappings.py @@ -1,6 +1,6 @@ from enum import IntEnum -from ..code_mappings import RoborockModeEnum +from ..code_mappings import RoborockEnum, RoborockModeEnum class B01_Q10_DP(RoborockModeEnum): @@ -354,3 +354,12 @@ class RemoteCommand(IntEnum): RIGHT = 3 STOP = 4 EXIT = 5 + + +class Q10RestrictionType(RoborockEnum): + """Restriction types supported by the Q10 map editor.""" + + unknown = -1 + NO_GO = 0 + NO_MOP = 2 + THRESHOLD = 3 diff --git a/roborock/devices/traits/b01/q10/map.py b/roborock/devices/traits/b01/q10/map.py index a33c2df22..cf4dbde5d 100644 --- a/roborock/devices/traits/b01/q10/map.py +++ b/roborock/devices/traits/b01/q10/map.py @@ -14,11 +14,12 @@ """ import logging +from collections.abc import Sequence from dataclasses import dataclass, field from typing import Any from roborock.data import RoborockBase -from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP, YXDeviceState +from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP, Q10RestrictionType, YXDeviceState from roborock.data.b01_q10.b01_q10_containers import Q10RoborockPoint from roborock.devices.traits.common import DpsDataConverter, TraitUpdateListener from roborock.exceptions import RoborockException @@ -29,8 +30,9 @@ Q10Room, Q10TracePacket, ) -from roborock.map.b01_q10_overlays import parse_virtual_wall_blob, parse_zone_blob +from roborock.map.b01_q10_overlays import Q10RestrictedZone, Q10VirtualWall, parse_virtual_wall_blob, parse_zone_blob from roborock.map.b01_q10_render import Q10MapOverlays, render_q10_map +from roborock.protocols.b01_q10_protocol import encode_restricted_zones, encode_room_name, encode_virtual_walls from .command import CommandTrait from .common import UpdatableTrait @@ -40,6 +42,10 @@ _DOCKED_STATES = {YXDeviceState.CHARGING, YXDeviceState.EMPTYING_THE_BIN} +def _to_roborock_point(point: tuple[int, int]) -> Q10RoborockPoint: + return Q10RoborockPoint.from_vector(*point) + + @dataclass class MapDps(RoborockBase): """Low-level map values delivered in the Q10 DPS stream.""" @@ -131,6 +137,64 @@ def path(self) -> list[Q10Point]: """Full path in the Q10 trace coordinate space used by the map renderer.""" return self._trace_packet.points if self._trace_packet else [] + @property + def restricted_zones(self) -> tuple[Q10RestrictedZone, ...]: + """Restrictions reported by the device in common Roborock coordinates.""" + restrictions: list[Q10RestrictedZone] = [] + for zone in self._map_dps.overlays.zones: + if len(zone.vertices) != 4: + continue + vertices = ( + _to_roborock_point(zone.vertices[0]), + _to_roborock_point(zone.vertices[1]), + _to_roborock_point(zone.vertices[2]), + _to_roborock_point(zone.vertices[3]), + ) + restriction_type = Q10RestrictionType(zone.type) + restrictions.append(Q10RestrictedZone(restriction_type, vertices)) + return tuple(restrictions) + + @property + def virtual_walls(self) -> tuple[Q10VirtualWall, ...]: + """Virtual walls reported by the device in common Roborock coordinates.""" + walls: list[Q10VirtualWall] = [] + for wall in self._map_dps.overlays.virtual_walls: + if len(wall.vertices) != 2: + continue + walls.append( + Q10VirtualWall( + start=_to_roborock_point(wall.vertices[0]), + end=_to_roborock_point(wall.vertices[1]), + ) + ) + return tuple(walls) + + async def set_restricted_zones(self, zones: Sequence[Q10RestrictedZone]) -> None: + """Replace the complete restricted-zone collection.""" + payload = encode_restricted_zones(zones) + await self._command.send( + B01_Q10_DP.COMMON, + {str(B01_Q10_DP.RESTRICTED_ZONE.code): payload}, + ) + + async def set_virtual_walls(self, walls: Sequence[Q10VirtualWall]) -> None: + """Replace the complete virtual-wall collection.""" + payload = encode_virtual_walls(walls) + await self._command.send( + B01_Q10_DP.COMMON, + {str(B01_Q10_DP.VIRTUAL_WALL.code): payload}, + ) + + async def set_room_name(self, room_id: int, name: str) -> None: + """Set the name of a room in the current map.""" + if room_id not in {room.id for room in self.rooms}: + raise RoborockException(f"Unknown Q10 room ID: {room_id}") + payload = encode_room_name(room_id, name) + await self._command.send( + B01_Q10_DP.COMMON, + {str(B01_Q10_DP.RESET_ROOM_NAME.code): payload}, + ) + @property def robot_position(self) -> Q10RoborockPoint | None: """Current position in the common Roborock millimetre coordinate space.""" diff --git a/roborock/devices/traits/b01/q10/maps.py b/roborock/devices/traits/b01/q10/maps.py index d09457764..a158c67ef 100644 --- a/roborock/devices/traits/b01/q10/maps.py +++ b/roborock/devices/traits/b01/q10/maps.py @@ -6,8 +6,9 @@ from roborock.data import RoborockBase from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP -from roborock.data.b01_q10.b01_q10_containers import dpMultiMap +from roborock.data.b01_q10.b01_q10_containers import Q10MapInfo, dpMultiMap from roborock.devices.traits.common import DpsDataConverter +from roborock.exceptions import RoborockException from .command import CommandTrait from .common import UpdatableTrait @@ -28,6 +29,13 @@ def current_map_id(self) -> str | None: return None return self.multi_map.current_map_id + @property + def map_list(self) -> list[Q10MapInfo]: + """Return a copy of the successfully reported saved-map list.""" + if self.multi_map is None or self.multi_map.op != "list" or self.multi_map.result != 1: + return [] + return list(self.multi_map.data) + class MapsTrait(Maps, UpdatableTrait): """Request and store the Q10 saved-map list.""" @@ -48,6 +56,20 @@ async def refresh(self) -> None: {str(B01_Q10_DP.MULTI_MAP.code): {"op": "list"}}, ) + async def set_current_map(self, map_id: str) -> None: + """Request a saved map without optimistically changing cached state.""" + if map_id not in {map_info.id for map_info in self.map_list}: + raise RoborockException(f"Unknown Q10 saved-map ID: {map_id}") + await self._command.send( + B01_Q10_DP.COMMON, + { + str(B01_Q10_DP.MULTI_MAP.code): { + "op": "apply", + "id": map_id, + } + }, + ) + def update_from_dps(self, decoded_dps: dict[B01_Q10_DP, Any]) -> None: """Store a successful saved-map list response.""" response = decoded_dps.get(B01_Q10_DP.MULTI_MAP) diff --git a/roborock/map/b01_q10_overlays.py b/roborock/map/b01_q10_overlays.py index 307ab210a..d55d328d7 100644 --- a/roborock/map/b01_q10_overlays.py +++ b/roborock/map/b01_q10_overlays.py @@ -31,6 +31,10 @@ import binascii from dataclasses import dataclass, field +from roborock.data import RoborockBase +from roborock.data.b01_q10.b01_q10_code_mappings import Q10RestrictionType +from roborock.data.b01_q10.b01_q10_containers import Q10RoborockPoint + _DEFAULT_RECORD_SIZE = 38 # 2-byte record header + up to 9 (x, y) int16 pairs @@ -42,6 +46,22 @@ class Q10Zone: vertices: list[tuple[int, int]] = field(default_factory=list) +@dataclass +class Q10RestrictedZone(RoborockBase): + """A four-point restriction in common Roborock coordinates.""" + + type: Q10RestrictionType + vertices: tuple[Q10RoborockPoint, Q10RoborockPoint, Q10RoborockPoint, Q10RoborockPoint] + + +@dataclass +class Q10VirtualWall(RoborockBase): + """A virtual wall in common Roborock coordinates.""" + + start: Q10RoborockPoint + end: Q10RoborockPoint + + def _decode_blob(data: str | None) -> bytes: if data is None: return b"" @@ -55,8 +75,8 @@ def parse_zone_blob(data: str | None) -> list[Q10Zone]: """Decode a Q10 zone/wall overlay blob into a list of :class:`Q10Zone`. Accepts the base64 string straight from the data point. Returns ``[]`` for - empty/absent/unparsable blobs (the device sends a single ``0x00`` byte when - there are none). + empty/absent/unparsable blobs. Empty DP 55 reports contain version ``1`` + and count ``0``; some firmware appends one padding byte. """ raw = _decode_blob(data) if len(raw) < 2: @@ -156,7 +176,7 @@ def parse_virtual_wall_blob(data: str | None) -> list[Q10Zone]: # Corrected from an earlier reading that treated type 3 as no-mop -- 3 is the # door-threshold rectangle; the no-mop area reads back as type 2. Reported and # verified by @andrewlyeats (ss07 read-backs + the ioBroker Q10 parser). -ZONE_TYPE_NO_GO = 0 +ZONE_TYPE_NO_GO = Q10RestrictionType.NO_GO.value ZONE_TYPE_VIRTUAL_WALL = 1 -ZONE_TYPE_NO_MOP = 2 -ZONE_TYPE_THRESHOLD = 3 +ZONE_TYPE_NO_MOP = Q10RestrictionType.NO_MOP.value +ZONE_TYPE_THRESHOLD = Q10RestrictionType.THRESHOLD.value diff --git a/roborock/protocols/b01_q10_protocol.py b/roborock/protocols/b01_q10_protocol.py index a3a18fa1b..4178cb1e0 100644 --- a/roborock/protocols/b01_q10_protocol.py +++ b/roborock/protocols/b01_q10_protocol.py @@ -1,13 +1,16 @@ """Roborock B01 Protocol encoding and decoding.""" +import base64 import json import logging +import struct from base64 import b64encode +from collections.abc import Sequence from dataclasses import dataclass from struct import pack from typing import Any -from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP +from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP, Q10RestrictionType from roborock.data.b01_q10.b01_q10_containers import Q10RoborockPoint from roborock.exceptions import RoborockException from roborock.map.b01_q10_map_parser import ( @@ -19,6 +22,7 @@ parse_map_packet, parse_trace_packet, ) +from roborock.map.b01_q10_overlays import Q10RestrictedZone, Q10VirtualWall from roborock.roborock_message import ( RoborockMessage, RoborockMessageProtocol, @@ -75,6 +79,118 @@ def encode_clean_params(params: CleanParams) -> str: return b64encode(payload).decode() +_Q10_ZONE_RECORD_SIZE = 38 +_Q10_ROOM_NAME_FIELD_SIZE = 19 + + +def _orientation(a: tuple[int, int], b: tuple[int, int], c: tuple[int, int]) -> int: + return (b[0] - a[0]) * (c[1] - a[1]) - (b[1] - a[1]) * (c[0] - a[0]) + + +def _on_segment(a: tuple[int, int], b: tuple[int, int], point: tuple[int, int]) -> bool: + return min(a[0], b[0]) <= point[0] <= max(a[0], b[0]) and min(a[1], b[1]) <= point[1] <= max(a[1], b[1]) + + +def _segments_cross(a: tuple[int, int], b: tuple[int, int], c: tuple[int, int], d: tuple[int, int]) -> bool: + orientations = ( + _orientation(a, b, c), + _orientation(a, b, d), + _orientation(c, d, a), + _orientation(c, d, b), + ) + if orientations[0] * orientations[1] < 0 and orientations[2] * orientations[3] < 0: + return True + return any( + orientation == 0 and _on_segment(start, end, point) + for orientation, start, end, point in ( + (orientations[0], a, b, c), + (orientations[1], a, b, d), + (orientations[2], c, d, a), + (orientations[3], c, d, b), + ) + ) + + +def _validate_quadrilateral(vertices: list[tuple[int, int]]) -> None: + area_twice = sum( + x1 * y2 - x2 * y1 for (x1, y1), (x2, y2) in zip(vertices, (*vertices[1:], vertices[0]), strict=True) + ) + if area_twice == 0: + raise ValueError("restricted-zone vertices must enclose an area") + if _segments_cross(vertices[0], vertices[1], vertices[2], vertices[3]) or _segments_cross( + vertices[1], vertices[2], vertices[3], vertices[0] + ): + raise ValueError("restricted-zone edges must not cross") + + +def encode_restricted_zones(zones: Sequence[Q10RestrictedZone]) -> str: + """Encode the complete Q10 restricted-zone collection.""" + if len(zones) > 255: + raise ValueError("zones must be a sequence of at most 255 restrictions") + if not zones: + return base64.b64encode(b"\x01\x00").decode() + + payload = bytearray((1, len(zones))) + for zone in zones: + if zone.type is Q10RestrictionType.unknown: + raise ValueError("unknown restricted-zone types cannot be written") + if len(zone.vertices) != 4: + raise ValueError("restricted zones must contain exactly four vertices") + if len(set(zone.vertices)) != 4: + raise ValueError("restricted-zone vertices must be distinct") + record = bytearray((zone.type, len(zone.vertices))) + vector_vertices = [point.to_vector() for point in zone.vertices] + _validate_quadrilateral(vector_vertices) + for x, y in vector_vertices: + record.extend( + struct.pack( + ">hh", + x, + y, + ) + ) + record.extend(bytes(_Q10_ZONE_RECORD_SIZE - len(record))) + payload.extend(record) + return base64.b64encode(payload).decode() + + +def encode_virtual_walls(walls: Sequence[Q10VirtualWall]) -> str: + """Encode the complete Q10 virtual-wall collection.""" + if len(walls) > 255: + raise ValueError("walls must be a sequence of at most 255 virtual walls") + + payload = bytearray((len(walls),)) + for wall in walls: + if wall.start == wall.end: + raise ValueError("virtual-wall endpoints must be distinct") + for point in (wall.start, wall.end): + x, y = point.to_vector() + payload.extend( + struct.pack( + ">hh", + x, + y, + ) + ) + return base64.b64encode(payload).decode() + + +def encode_room_name(room_id: int, name: str) -> str: + """Encode a Q10 room-name update.""" + if not 0 <= room_id <= 255: + raise ValueError("room_id must be an integer between 0 and 255") + if not name or "\x00" in name: + raise ValueError("name must be a non-empty string without null bytes") + encoded_name = name.encode("utf-8") + if len(encoded_name) > _Q10_ROOM_NAME_FIELD_SIZE: + raise ValueError("name must be at most 19 UTF-8 bytes") + + payload = bytearray((1, room_id, 0, len(encoded_name))) + payload.extend(encoded_name) + payload.extend(bytes(_Q10_ROOM_NAME_FIELD_SIZE - len(encoded_name))) + return base64.b64encode(payload).decode() + + def encode_mqtt_payload(command: B01_Q10_DP, params: ParamsType) -> RoborockMessage: """Encode payload for B01 Q10 commands over MQTT. diff --git a/tests/devices/traits/b01/q10/test_map.py b/tests/devices/traits/b01/q10/test_map.py index 060553af5..1f6faf3f2 100644 --- a/tests/devices/traits/b01/q10/test_map.py +++ b/tests/devices/traits/b01/q10/test_map.py @@ -16,7 +16,8 @@ import pytest from roborock.cli import _await_q10_map_push, cli -from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP, YXDeviceState +from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP, Q10RestrictionType, YXDeviceState +from roborock.data.b01_q10.b01_q10_containers import Q10RoborockPoint from roborock.devices.traits.b01.q10 import Q10PropertiesApi, create from roborock.devices.traits.b01.q10.command import CommandTrait from roborock.devices.traits.b01.q10.map import MapContentTrait, MapDpsTrait @@ -28,6 +29,7 @@ parse_map_packet, parse_trace_packet, ) +from roborock.map.b01_q10_overlays import Q10RestrictedZone, Q10VirtualWall from roborock.map.b01_q10_render import Q10MapOverlays from roborock.protocols.b01_q10_protocol import Q10DpsUpdate, Q10Message @@ -586,3 +588,177 @@ def test_map_content_trait_as_dict_camelizes_child_keys() -> None: } assert data["path"] == [{"x": 100, "y": 200}, {"x": 150, "y": 250}] assert data["robotPosition"] == {"x": 25875, "y": 26125} + + +async def test_set_current_map_uses_apply_for_listed_string_id( + q10_api: Q10PropertiesApi, mock_channel: FakeB01Q10Channel +) -> None: + q10_api.maps.update_from_dps( + { + B01_Q10_DP.MULTI_MAP: { + "data": [{"id": "12345"}, {"id": "67890"}], + "op": "list", + "result": 1, + } + } + ) + + await q10_api.maps.set_current_map("67890") + + assert mock_channel.published_commands == [ + ( + B01_Q10_DP.COMMON, + {str(B01_Q10_DP.MULTI_MAP.code): {"op": "apply", "id": "67890"}}, + ) + ] + assert q10_api.maps.current_map_id == "12345" + + +@pytest.mark.parametrize("map_id", ["999"]) +async def test_set_current_map_rejects_unknown_or_non_string_id( + q10_api: Q10PropertiesApi, + map_id: str, + mock_channel: FakeB01Q10Channel, +) -> None: + q10_api.maps.update_from_dps( + { + B01_Q10_DP.MULTI_MAP: { + "data": [{"id": "12345"}], + "op": "list", + "result": 1, + } + } + ) + + with pytest.raises(RoborockException, match="Unknown Q10 saved-map ID"): + await q10_api.maps.set_current_map(map_id) + + assert mock_channel.published_commands == [] + + +def test_map_content_exposes_restrictions_in_common_coordinates() -> None: + map_dps = MapDpsTrait() + trait = _map_trait(map_dps) + + map_dps.update_from_dps( + { + B01_Q10_DP.RESTRICTED_ZONE_UP: _zone_blob(), + B01_Q10_DP.VIRTUAL_WALL_UP: "AQAAAAAACgAU", + } + ) + + assert trait.restricted_zones == ( + Q10RestrictedZone( + Q10RestrictionType.NO_GO, + ( + Q10RoborockPoint(25500, 25500), + Q10RoborockPoint(25700, 25500), + Q10RoborockPoint(25700, 25700), + Q10RoborockPoint(25500, 25700), + ), + ), + ) + assert trait.virtual_walls == ( + Q10VirtualWall( + start=Q10RoborockPoint(25500, 25500), + end=Q10RoborockPoint(25550, 25600), + ), + ) + + +def test_map_content_preserves_unknown_restriction_type() -> None: + map_dps = MapDpsTrait() + trait = _map_trait(map_dps) + raw = bytearray(base64.b64decode(_zone_blob())) + raw[2] = 42 + + map_dps.update_from_dps({B01_Q10_DP.RESTRICTED_ZONE_UP: base64.b64encode(raw).decode()}) + + assert trait.restricted_zones[0].type is Q10RestrictionType.unknown + + +async def test_set_restricted_zones_sends_atomic_common_write( + q10_api: Q10PropertiesApi, mock_channel: FakeB01Q10Channel +) -> None: + zone = Q10RestrictedZone( + Q10RestrictionType.NO_MOP, + ( + Q10RoborockPoint(25500, 25500), + Q10RoborockPoint(25550, 25500), + Q10RoborockPoint(25550, 25550), + Q10RoborockPoint(25500, 25550), + ), + ) + + await q10_api.map.set_restricted_zones([zone]) + + assert mock_channel.published_commands == [ + ( + B01_Q10_DP.COMMON, + {str(B01_Q10_DP.RESTRICTED_ZONE.code): "AQECBAAAAAAACgAAAAoACgAAAAoAAAAAAAAAAAAAAAAAAAAAAAAAAA=="}, + ) + ] + assert q10_api.map.restricted_zones == () + + +async def test_set_virtual_walls_sends_atomic_common_write( + q10_api: Q10PropertiesApi, mock_channel: FakeB01Q10Channel +) -> None: + wall = Q10VirtualWall( + start=Q10RoborockPoint(25450, 25500), + end=Q10RoborockPoint(25550, 25600), + ) + + await q10_api.map.set_virtual_walls((wall,)) + + assert mock_channel.published_commands == [ + ( + B01_Q10_DP.COMMON, + {str(B01_Q10_DP.VIRTUAL_WALL.code): "Af/2AAAACgAU"}, + ) + ] + assert q10_api.map.virtual_walls == () + + +async def test_invalid_map_control_does_not_publish(q10_api: Q10PropertiesApi, mock_channel: FakeB01Q10Channel) -> None: + invalid = Q10VirtualWall( + start=Q10RoborockPoint(25501, 25500), + end=Q10RoborockPoint(25550, 25600), + ) + + with pytest.raises(ValueError, match="5 mm grid"): + await q10_api.map.set_virtual_walls([invalid]) + + assert mock_channel.published_commands == [] + + +async def test_set_room_name_sends_validated_common_write( + q10_api: Q10PropertiesApi, mock_channel: FakeB01Q10Channel +) -> None: + q10_api.map.update_from_map_packet(parse_map_packet(FIXTURE.read_bytes())) + + await q10_api.map.set_room_name(2, "Study") + + assert mock_channel.published_commands == [ + ( + B01_Q10_DP.COMMON, + {str(B01_Q10_DP.RESET_ROOM_NAME.code): "AQIABVN0dWR5AAAAAAAAAAAAAAAAAAA="}, + ) + ] + assert q10_api.map.rooms[0].raw_name == "rr_living_room" + + +async def test_set_room_name_rejects_unknown_room_before_publish( + q10_api: Q10PropertiesApi, mock_channel: FakeB01Q10Channel +) -> None: + q10_api.map.update_from_map_packet(parse_map_packet(FIXTURE.read_bytes())) + + with pytest.raises(RoborockException, match="Unknown Q10 room ID"): + await q10_api.map.set_room_name(99, "Study") + + assert mock_channel.published_commands == [] + + +async def test_clear_restrictions_without_snapshot(q10_api: Q10PropertiesApi, mock_channel: FakeB01Q10Channel) -> None: + await q10_api.map.set_restricted_zones([]) + assert mock_channel.published_commands == [(B01_Q10_DP.COMMON, {"54": "AQA="})] diff --git a/tests/protocols/test_b01_q10_protocol.py b/tests/protocols/test_b01_q10_protocol.py index 16fe9a108..8fd745e09 100644 --- a/tests/protocols/test_b01_q10_protocol.py +++ b/tests/protocols/test_b01_q10_protocol.py @@ -1,5 +1,6 @@ """Tests for the B01 protocol message encoding and decoding.""" +import base64 import json import logging import pathlib @@ -11,11 +12,12 @@ from freezegun import freeze_time from syrupy import SnapshotAssertion -from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP, YXWaterLevel +from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP, Q10RestrictionType, YXWaterLevel from roborock.data.b01_q10.b01_q10_containers import Q10RoborockPoint from roborock.data.code_mappings import completed_warnings from roborock.exceptions import RoborockException from roborock.map.b01_q10_map_parser import Q10CleanRecordDetail, Q10MapPacket, Q10MapPacketKind, Q10TracePacket +from roborock.map.b01_q10_overlays import Q10RestrictedZone, Q10VirtualWall from roborock.protocols.b01_q10_protocol import ( CleanParams, Q10DpsUpdate, @@ -23,6 +25,9 @@ decode_rpc_response, encode_clean_params, encode_mqtt_payload, + encode_restricted_zones, + encode_room_name, + encode_virtual_walls, ) from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol @@ -289,3 +294,185 @@ def test_encode_mqtt_payload(command: B01_Q10_DP, params: dict[str, Any], snapsh assert snapshot == message.payload json.loads(message.payload.decode()) + + +def _point(x: int, y: int) -> Q10RoborockPoint: + """Build a public-coordinate point from a captured Q10 wire vector.""" + return Q10RoborockPoint.from_vector(x, y) + + +def _zone(zone_type: Q10RestrictionType, *vertices: tuple[int, int]) -> Q10RestrictedZone: + return Q10RestrictedZone(zone_type, tuple(_point(*vertex) for vertex in vertices)) # type: ignore[arg-type] + + +def _wall(start: tuple[int, int], end: tuple[int, int]) -> Q10VirtualWall: + return Q10VirtualWall(_point(*start), _point(*end)) + + +def test_encode_restricted_zones_real_three_zone_capture_byte_exact() -> None: + """Encoding is the exact inverse of a real three-zone DP 55 capture.""" + captured = ( + "AQMABP9A9ToAEPU6ABDzzv9A884AAAAAAAAAAAAAAAAAAAAAAAAAAAAE/Rb/mv5z/5r+c/3L/Rb9ywAAAA" + "AAAAAAAAAAAAAAAAAAAAAAAAQDpADEB3UAxAd1/GMDpPxjAAAAAAAAAAAAAAAAAAAAAAAAAAA=" + ) + zones = [ + _zone(Q10RestrictionType.NO_GO, (-192, -2758), (16, -2758), (16, -3122), (-192, -3122)), + _zone(Q10RestrictionType.NO_GO, (-746, -102), (-397, -102), (-397, -565), (-746, -565)), + _zone(Q10RestrictionType.NO_GO, (932, 196), (1909, 196), (1909, -925), (932, -925)), + ] + + assert encode_restricted_zones(zones) == captured + + +def test_encode_restricted_zones_empty_is_device_canonical_payload() -> None: + assert encode_restricted_zones([]) == "AQA=" + + +def test_encode_restricted_zones_mixed_types_and_padding_byte_exact() -> None: + """Multiple restriction kinds retain order and fixed 38-byte padding.""" + zones = [ + _zone(Q10RestrictionType.NO_MOP, (-1, 2), (3, 2), (3, 8), (-1, 8)), + _zone(Q10RestrictionType.THRESHOLD, (-32_768, 32_767), (0, 32_767), (0, 5), (-32_768, 5)), + ] + record_one = bytes.fromhex("0204ffff00020003000200030008ffff0008") + bytes(20) + record_two = bytes.fromhex("030480007fff00007fff0000000580000005") + bytes(20) + + assert base64.b64decode(encode_restricted_zones(zones)) == b"\x01\x02" + record_one + record_two + + +def test_encode_virtual_walls_real_mixed_capture_byte_exact() -> None: + """Encoding preserves the proven x/y order for real horizontal and vertical walls.""" + captured = "AgNyBGMFsARjAf8FAAH7CnE=" + walls = [_wall((882, 1123), (1456, 1123)), _wall((511, 1280), (507, 2673))] + + assert encode_virtual_walls(walls) == captured + + +def test_encode_virtual_walls_empty_is_device_canonical_payload() -> None: + assert encode_virtual_walls([]) == "AA==" + + +def test_encode_virtual_walls_signed_boundaries_byte_exact() -> None: + walls = [_wall((-32_768, 32_767), (32_767, -32_768))] + + assert base64.b64decode(encode_virtual_walls(walls)) == bytes.fromhex("0180007fff7fff8000") + + +def test_encode_restricted_zones_rejects_more_than_wire_count() -> None: + zone = _zone(Q10RestrictionType.NO_GO, (0, 0), (1, 0), (1, 1), (0, 1)) + + with pytest.raises(ValueError, match="at most 255"): + encode_restricted_zones([zone] * 256) + + +def test_encode_restricted_zones_rejects_unknown_wire_type() -> None: + zone = Q10RestrictedZone(Q10RestrictionType.unknown, (_point(0, 0), _point(1, 0), _point(1, 1), _point(0, 1))) + + with pytest.raises(ValueError, match="unknown"): + encode_restricted_zones([zone]) + + +@pytest.mark.parametrize( + "vertices", + [ + ((0, 0), (1, 0), (1, 1)), + ((0, 0), (1, 0), (1, 1), (0, 1), (-1, 0)), + ], +) +def test_encode_restricted_zones_rejects_wrong_cardinality(vertices: tuple[tuple[int, int], ...]) -> None: + zone = Q10RestrictedZone( + Q10RestrictionType.NO_GO, + tuple(_point(*vertex) for vertex in vertices), # type: ignore[arg-type] + ) + + with pytest.raises(ValueError, match="exactly four"): + encode_restricted_zones([zone]) + + +def test_encode_restricted_zones_rejects_duplicate_vertices() -> None: + duplicate = _point(0, 0) + zone = Q10RestrictedZone( + Q10RestrictionType.NO_GO, + (duplicate, _point(1, 0), _point(1, 1), duplicate), + ) + + with pytest.raises(ValueError, match="distinct"): + encode_restricted_zones([zone]) + + +@pytest.mark.parametrize( + ("vertices", "message"), + [ + (((0, 0), (1, 0), (2, 0), (3, 0)), "enclose an area"), + (((0, 0), (2, 2), (0, 2), (2, 0)), "enclose an area|must not cross"), + ], +) +def test_encode_restricted_zones_rejects_invalid_geometry( + vertices: tuple[tuple[int, int], ...], + message: str, +) -> None: + zone = _zone(Q10RestrictionType.NO_GO, *vertices) + + with pytest.raises(ValueError, match=message): + encode_restricted_zones([zone]) + + +def test_encode_virtual_walls_rejects_more_than_wire_count() -> None: + wall = _wall((0, 0), (1, 1)) + + with pytest.raises(ValueError, match="at most 255"): + encode_virtual_walls([wall] * 256) + + +def test_encode_virtual_walls_rejects_duplicate_endpoints() -> None: + point = _point(0, 0) + + with pytest.raises(ValueError, match="distinct"): + encode_virtual_walls([Q10VirtualWall(point, point)]) + + +@pytest.mark.parametrize( + "point", + [ + Q10RoborockPoint(25_501, 25_500), + Q10RoborockPoint(-138_345, 25_500), + ], +) +def test_zone_and_wall_encoders_reject_invalid_coordinates(point: Q10RoborockPoint) -> None: + valid = (_point(10, 10), _point(11, 10), _point(11, 11)) + zone = Q10RestrictedZone(Q10RestrictionType.NO_GO, (point, *valid)) + wall = Q10VirtualWall(point, _point(2, 2)) + + with pytest.raises(ValueError, match="coordinates"): + encode_restricted_zones([zone]) + with pytest.raises(ValueError, match="coordinates"): + encode_virtual_walls([wall]) + + +def test_overlay_encoders_do_not_mutate_callers() -> None: + zones = [_zone(Q10RestrictionType.NO_MOP, (0, 0), (2, 0), (2, 2), (0, 2))] + walls = [_wall((-1, -1), (1, 1))] + original_zones = list(zones) + original_walls = list(walls) + + encode_restricted_zones(zones) + encode_virtual_walls(walls) + + assert zones == original_zones + assert walls == original_walls + + +def test_encode_room_name_uses_fixed_utf8_field() -> None: + assert base64.b64decode(encode_room_name(7, "Café")) == bytes((1, 7, 0, 5)) + "Café".encode() + bytes(14) + + +@pytest.mark.parametrize("room_id", [-1, 256]) +def test_encode_room_name_rejects_invalid_room_id(room_id: int) -> None: + with pytest.raises(ValueError, match="room_id"): + encode_room_name(room_id, "Study") + + +@pytest.mark.parametrize("name", ["", "a\x00b", "x" * 20, "é" * 10]) +def test_encode_room_name_rejects_invalid_name(name: str) -> None: + with pytest.raises(ValueError, match="name"): + encode_room_name(7, name)