test: add end-to-end multi-part map content trait tests

roborock/testing/v1_simulator.py

@@ -10,7 +10,7 @@
 import logging
 import struct
 import time
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from dataclasses import asdict, replace
 from enum import Enum
 from typing import Any
@@ -45,6 +45,7 @@
 from roborock.roborock_message import RoborockDataProtocol, RoborockMessage, RoborockMessageProtocol
 from roborock.testing.channel import FakeChannel
 from roborock.testing.simulator import RoborockDeviceSimulator
+from roborock.util import get_next_int
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -255,8 +256,12 @@ def __init__(
 
         # The simulated robot's own view of its map. Tests extend the path
         # with advance_path() and move the robot via robot_position.
+        # Set map_parts > 1 (and optionally map_part_order) to split map
+        # responses across MAP_RESPONSE_PART (300) and MAP_RESPONSE (301) frames.
         self.path_length = 0
         self.robot_position = (25000, 24000, 90)
+        self.map_parts = 1
+        self.map_part_order: Sequence[int] | None = None
 
         self.device_cache = DeviceCache(self.duid, InMemoryCache())
         # The endpoint must fit the eight byte map response header field and the
@@ -416,11 +421,44 @@ def build_map_bytes(self) -> bytes:
         data = header + body + digest_header
         return data + hashlib.sha1(data, usedforsecurity=False).digest()
 
-    def _map_response_message(self, request_id: int, map_bytes: bytes) -> RoborockMessage:
-        """Frame the map as the encrypted binary response the device sends."""
+    def _map_response_messages(self, request_id: int, map_bytes: bytes) -> list[RoborockMessage]:
+        """Frame the map as one or more encrypted binary messages as the device sends."""
+        if self.map_parts < 1:
+            raise ValueError("map_parts must be at least 1")
         body = Utils.encrypt_cbc(gzip.compress(map_bytes), self.security_data.nonce)
         header = struct.pack("<8s8sH6s", self.security_data.endpoint.encode(), bytes(8), request_id & 0xFFFF, bytes(6))
-        return RoborockMessage(protocol=RoborockMessageProtocol.MAP_RESPONSE, payload=header + body)
+        if self.map_parts == 1:
+            if self.map_part_order is not None and list(self.map_part_order) != [0]:
+                raise ValueError("map_part_order must be a permutation of range(map_parts)")
+            return [RoborockMessage(protocol=RoborockMessageProtocol.MAP_RESPONSE, payload=header + body)]
+        if self.map_parts > len(body):
+            raise ValueError("map_parts cannot exceed encrypted map body length")
+
+        # Keep the 24-byte header intact on the first part and split the encrypted
+        # body evenly into map_parts non-empty contiguous slices.
+        total = len(body)
+        boundaries = [(i * total) // self.map_parts for i in range(self.map_parts + 1)]
+        chunks = [body[boundaries[i] : boundaries[i + 1]] for i in range(self.map_parts)]
+        chunks[0] = header + chunks[0]
+        random_id = get_next_int(10000, 99999)
+        timestamp = int(time.time())
+        messages = [
+            RoborockMessage(
+                protocol=RoborockMessageProtocol.MAP_RESPONSE
+                if idx == len(chunks) - 1
+                else RoborockMessageProtocol.MAP_RESPONSE_PART,
+                payload=chunk,
+                seq=idx + 1,
+                random=random_id,
+                timestamp=timestamp,
+            )
+            for idx, chunk in enumerate(chunks)
+        ]
+        if self.map_part_order is not None:
+            if sorted(self.map_part_order) != list(range(self.map_parts)):
+                raise ValueError("map_part_order must be a permutation of range(map_parts)")
+            messages = [messages[i] for i in self.map_part_order]
+        return messages
 
     async def _handle_publish(self, message: RoborockMessage, channel: FakeChannel) -> None:
         if not message.payload:
@@ -455,7 +493,8 @@ async def _handle_publish(self, message: RoborockMessage, channel: FakeChannel)
 
         if method == "get_map_v1" and error is None and isinstance(result, bytes):
             # Custom handlers return decompressed map bytes, just like the default.
-            channel.notify_subscribers(self._map_response_message(msg_id, result))
+            for map_msg in self._map_response_messages(msg_id, result):
+                channel.notify_subscribers(map_msg)
             return
 
         response_data = {

tests/devices/traits/v1/test_map_content.py

@@ -1,87 +1,192 @@
-"""Tests for the MapContentTrait."""
+"""End-to-end tests for MapContentTrait backed by V1VacuumSimulator."""
 
-from unittest.mock import AsyncMock, MagicMock, patch
+from collections.abc import AsyncGenerator
+from dataclasses import replace
 
 import pytest
 
+from roborock.data import HomeDataDevice, HomeDataProduct, RoborockDockTypeCode
 from roborock.devices.device import RoborockDevice
+from roborock.devices.rpc.v1_channel import V1Channel
 from roborock.devices.traits.v1.map_content import MapContentTrait
-from roborock.map.map_parser import ParsedMapData
-from roborock.roborock_typing import RoborockCommand
+from roborock.exceptions import RoborockException
+from roborock.protocols.v1_protocol import V1RpcChannel
+from roborock.testing import DEFAULT_STATUS, V1VacuumSimulator
+
+
+@pytest.fixture(name="simulator")
+def simulator_fixture(
+    device_info: HomeDataDevice,
+    products: list[HomeDataProduct],
+    dock_type_code: RoborockDockTypeCode | None,
+) -> V1VacuumSimulator:
+    """Create a connected V1VacuumSimulator backing the device."""
+    product = next(p for p in products if p.id == device_info.product_id)
+    sim = V1VacuumSimulator(
+        duid=device_info.duid,
+        device_info=device_info,
+        product=product,
+        status=replace(DEFAULT_STATUS, dock_type=dock_type_code),
+    )
+    sim.connect()
+    return sim
+
+
+@pytest.fixture(name="channel")
+def simulator_channel_fixture(simulator: V1VacuumSimulator) -> V1Channel:
+    """Provide the simulator's real V1Channel."""
+    return simulator.v1_channel
+
+
+@pytest.fixture(name="mock_rpc_channel")
+def simulator_rpc_channel_fixture(simulator: V1VacuumSimulator) -> V1RpcChannel:
+    """Provide the simulator's real combined RPC channel."""
+    return simulator.v1_channel.rpc_channel
+
+
+@pytest.fixture(name="mock_mqtt_rpc_channel")
+def simulator_mqtt_rpc_channel_fixture(simulator: V1VacuumSimulator) -> V1RpcChannel:
+    """Provide the simulator's real MQTT RPC channel."""
+    return simulator.v1_channel.mqtt_rpc_channel
+
+
+@pytest.fixture(name="mock_map_rpc_channel")
+def simulator_map_rpc_channel_fixture(simulator: V1VacuumSimulator) -> V1RpcChannel:
+    """Provide the simulator's real map RPC channel."""
+    return simulator.v1_channel.map_rpc_channel
+
+
+@pytest.fixture(name="mock_blob_rpc_channel")
+def simulator_blob_rpc_channel_fixture(simulator: V1VacuumSimulator) -> V1RpcChannel:
+    """Provide the simulator's real blob RPC channel."""
+    return simulator.v1_channel.blob_rpc_channel
+
+
+@pytest.fixture(autouse=True)
+async def discover_features_fixture(device: RoborockDevice) -> AsyncGenerator[None, None]:
+    """Connect the device against the simulator and tear down cleanly."""
+    await device.connect()
+    try:
+        yield
+    finally:
+        await device.close()
 
 
-@pytest.fixture
-def map_content_trait(device: RoborockDevice) -> MapContentTrait:
-    """Create a MapContentTrait instance with mocked dependencies."""
+@pytest.fixture(name="map_content_trait")
+def map_content_trait_fixture(device: RoborockDevice) -> MapContentTrait:
+    """Return the device's MapContentTrait."""
     assert device.v1_properties
     return device.v1_properties.map_content
 
 
+@pytest.mark.parametrize(
+    ("map_parts", "map_part_order"),
+    [
+        (1, None),
+        (2, [0, 1]),
+        (2, [1, 0]),
+        (3, [0, 1, 2]),
+        (3, [2, 0, 1]),
+        (4, [3, 1, 2, 0]),
+    ],
+    ids=[
+        "single_part",
+        "two_parts_in_order",
+        "two_parts_out_of_order",
+        "three_parts_in_order",
+        "three_parts_out_of_order",
+        "four_parts_permuted",
+    ],
+)
 async def test_refresh_map_content_trait(
+    simulator: V1VacuumSimulator,
     map_content_trait: MapContentTrait,
-    mock_map_rpc_channel: AsyncMock,
+    map_parts: int,
+    map_part_order: list[int] | None,
 ) -> None:
-    """Test successfully getting and parsing map content."""
-    map_data = b"dummy_map_bytes"
-    mock_map_rpc_channel.send_command.return_value = map_data
-    mock_parsed_data = ParsedMapData(
-        image_content=b"dummy_image_content",
-        map_data=MagicMock(),
-    )
+    """Test refreshing and parsing single-part and multi-part map responses end-to-end."""
+    simulator.advance_path(80)
+    simulator.robot_position = (31000, 32000, 180)
+    simulator.map_parts = map_parts
+    simulator.map_part_order = map_part_order
 
-    with patch("roborock.devices.traits.v1.map_content.MapParser.parse", return_value=mock_parsed_data) as mock_parse:
-        await map_content_trait.refresh()
-        mock_parse.assert_called_once_with(map_data)
+    await map_content_trait.refresh()
 
-    assert map_content_trait.image_content == b"dummy_image_content"
+    assert map_content_trait.image_content is not None
+    assert map_content_trait.image_content.startswith(b"\x89PNG")
+    assert map_content_trait.raw_api_response == simulator.build_map_bytes()
     assert map_content_trait.map_data is not None
-
-    mock_map_rpc_channel.send_command.assert_called_once_with(RoborockCommand.GET_MAP_V1)
+    assert map_content_trait.map_data.path is not None
+    assert len(map_content_trait.map_data.path.path[0]) == 80
+    assert map_content_trait.map_data.vacuum_position is not None
+    assert (
+        map_content_trait.map_data.vacuum_position.x,
+        map_content_trait.map_data.vacuum_position.y,
+    ) == (31000, 32000)
 
 
 async def test_refresh_unchanged_content_skips_parsing(
+    simulator: V1VacuumSimulator,
     map_content_trait: MapContentTrait,
-    mock_map_rpc_channel: AsyncMock,
 ) -> None:
-    """Test that refresh avoids re-parsing when the map content is unchanged."""
-    map_data = b"dummy_map_bytes"
-    mock_map_rpc_channel.send_command.return_value = map_data
-    mock_parsed_data = ParsedMapData(
-        image_content=b"dummy_image_content",
-        map_data=MagicMock(),
-    )
+    """Test that refresh avoids re-parsing unchanged maps and updates when map bytes change."""
+    simulator.advance_path(40)
+    simulator.robot_position = (28000, 29000, 90)
+    simulator.map_parts = 3
+
+    await map_content_trait.refresh()
+    initial_image = map_content_trait.image_content
+    initial_map_data = map_content_trait.map_data
+    assert initial_image is not None
+    assert initial_map_data is not None
+
+    # Second refresh with identical multi-part map preserves the existing MapData instance
+    await map_content_trait.refresh()
+    assert map_content_trait.map_data is initial_map_data
+
+    # Third refresh after robot movement re-parses and renders the updated map
+    simulator.advance_path(40)
+    simulator.robot_position = (33000, 34000, 270)
+    simulator.map_part_order = [2, 0, 1]
+    await map_content_trait.refresh()
+
+    assert map_content_trait.map_data is not initial_map_data
+    assert map_content_trait.image_content != initial_image
+    assert map_content_trait.map_data is not None
+    assert map_content_trait.map_data.path is not None
+    assert len(map_content_trait.map_data.path.path[0]) == 80
+    assert map_content_trait.map_data.vacuum_position is not None
+    assert (
+        map_content_trait.map_data.vacuum_position.x,
+        map_content_trait.map_data.vacuum_position.y,
+    ) == (33000, 34000)
 
-    with patch("roborock.devices.traits.v1.map_content.MapParser.parse", return_value=mock_parsed_data) as mock_parse:
-        await map_content_trait.refresh()
-        assert mock_parse.call_count == 1
-        assert map_content_trait.image_content == b"dummy_image_content"
 
-        # Second refresh with identical response should not re-parse
-        await map_content_trait.refresh()
-        assert mock_parse.call_count == 1
-
-        # Third refresh with new response parses again
-        new_map_data = b"new_map_bytes"
-        new_parsed_data = ParsedMapData(
-            image_content=b"new_image_content",
-            map_data=MagicMock(),
-        )
-        mock_map_rpc_channel.send_command.return_value = new_map_data
-        mock_parse.return_value = new_parsed_data
+async def test_refresh_invalid_map_raises_exception(
+    simulator: V1VacuumSimulator,
+    map_content_trait: MapContentTrait,
+) -> None:
+    """Test malformed map payload raises RoborockException after multi-part reassembly."""
+    simulator.custom_handlers["get_map_v1"] = lambda _: b"corrupt_map_bytes"
+    simulator.map_parts = 2
+
+    with pytest.raises(RoborockException, match="Failed to parse map data"):
         await map_content_trait.refresh()
-        assert mock_parse.call_count == 2
-        assert map_content_trait.image_content == b"new_image_content"
 
 
-async def test_async_parse_map_content(map_content_trait: MapContentTrait) -> None:
-    """Test async_parse_map_content on MapContentConverter."""
-    map_data = b"dummy_map_bytes"
-    mock_parsed_data = ParsedMapData(
-        image_content=b"dummy_image_content",
-        map_data=MagicMock(),
-    )
-    with patch("roborock.devices.traits.v1.map_content.MapParser.parse", return_value=mock_parsed_data) as mock_parse:
-        result = await map_content_trait.converter.async_parse_map_content(map_data)
-        mock_parse.assert_called_once_with(map_data)
-        assert result.image_content == b"dummy_image_content"
-        assert result.raw_api_response == map_data
+async def test_async_parse_map_content(
+    simulator: V1VacuumSimulator,
+    map_content_trait: MapContentTrait,
+) -> None:
+    """Test async_parse_map_content on MapContentConverter with real map bytes."""
+    simulator.advance_path(25)
+    map_bytes = simulator.build_map_bytes()
+
+    result = await map_content_trait.converter.async_parse_map_content(map_bytes)
+
+    assert result.image_content is not None
+    assert result.image_content.startswith(b"\x89PNG")
+    assert result.raw_api_response == map_bytes
+    assert result.map_data is not None
+    assert result.map_data.path is not None
+    assert len(result.map_data.path.path[0]) == 25

tests/testing/test_v1_simulator.py

@@ -13,6 +13,7 @@
 from roborock.devices.traits.v1.consumeable import ConsumableAttribute
 from roborock.exceptions import RoborockException
 from roborock.map.map_parser import MapParser, MapParserConfig
+from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
 from roborock.testing import (
     DEFAULT_NETWORK_INFO,
     DEFAULT_STATUS,
@@ -308,6 +309,43 @@ async def test_full_map_refresh_from_simulator():
     assert (map_content.map_data.vacuum_position.x, map_content.map_data.vacuum_position.y) == (31000, 32000)
 
 
+async def test_simulator_multipart_map_framing() -> None:
+    """Configuring map_parts and map_part_order frames ordered MAP_RESPONSE_PART + MAP_RESPONSE chunks."""
+    simulator = V1VacuumSimulator(duid="s7_framing")
+    simulator.connect()
+    simulator.advance_path(40)
+    simulator.map_parts = 3
+    simulator.map_part_order = [2, 0, 1]
+
+    emitted: list[RoborockMessage] = []
+    unsub = await simulator.mqtt_channel.subscribe(emitted.append)
+    try:
+        raw = await simulator.v1_channel.map_rpc_channel.send_command("get_map_v1")
+    finally:
+        unsub()
+
+    assert [(msg.protocol, msg.seq) for msg in emitted] == [
+        (RoborockMessageProtocol.MAP_RESPONSE, 3),
+        (RoborockMessageProtocol.MAP_RESPONSE_PART, 1),
+        (RoborockMessageProtocol.MAP_RESPONSE_PART, 2),
+    ]
+    first_part = next(msg for msg in emitted if msg.seq == 1)
+    assert first_part.payload is not None
+    assert first_part.payload.startswith(simulator.security_data.endpoint.encode())
+    assert len(first_part.payload) > 24
+    assert raw == simulator.build_map_bytes()
+
+
+async def test_simulator_multipart_preserves_header_at_high_part_count() -> None:
+    """High part counts keep the complete 24-byte header on the first part."""
+    simulator = V1VacuumSimulator(duid="s7_many_parts")
+    simulator.connect()
+    simulator.map_parts = 20
+
+    raw = await simulator.v1_channel.map_rpc_channel.send_command("get_map_v1")
+    assert raw == simulator.build_map_bytes()
+
+
 @pytest.mark.parametrize("count", [0, 40, 400])
 def test_simulated_map_length_and_digest(count: int) -> None:
     """Synthetic maps satisfy the app's full-file length and SHA-1 checks."""