fix: fix room mapping parsing bug and add addtiional format samples

Author: allenporter

roborock/devices/traits/v1/rooms.py

@@ -75,17 +75,17 @@ def _extract_segment_pairs(response: list) -> list[tuple[int, str]]:
     list of lists, where each inner list is a pair of [segment_id, iot_id]. This
     function normalizes the response into a list of (segment_id, iot_id) tuples
 
-    NOTE: We currently only have samples of the list of lists format in
-    tests/protocols/testdata so improving test coverage with samples from a real
-    device with this format would be helpful.
+    NOTE: We currently only partial samples of the room mapping formats, so
+    improving test coverage with samples from a real device with this format
+    would be helpful.
     """
     if len(response) == 2 and not isinstance(response[0], list):
         segment_id, iot_id = response[0], response[1]
         return [(segment_id, iot_id)]
 
     segment_pairs: list[tuple[int, str]] = []
     for part in response:
-        if not isinstance(part, list) or len(part) != 2:
+        if not isinstance(part, list) or len(part) < 2:
             _LOGGER.warning("Unexpected room mapping entry format: %r", part)
             continue
         segment_id, iot_id = part[0], part[1]

tests/devices/traits/v1/test_rooms.py

@@ -3,18 +3,13 @@
 from unittest.mock import AsyncMock
 
 import pytest
+from typing import Any
 
 from roborock.devices.device import RoborockDevice
 from roborock.devices.traits.v1.rooms import RoomsTrait
 from roborock.devices.traits.v1.status import StatusTrait
 from roborock.roborock_typing import RoborockCommand
 
-# Rooms from mock_data.HOME_DATA
-# {"id": 2362048, "name": "Example room 1"},
-# {"id": 2362044, "name": "Example room 2"},
-# {"id": 2362041, "name": "Example room 3"},
-ROOM_MAPPING_DATA = [[16, "2362048"], [17, "2362044"], [18, "2362041"]]
-
 
 @pytest.fixture
 def status_trait(device: RoborockDevice) -> StatusTrait:
@@ -30,15 +25,25 @@ def rooms_trait(device: RoborockDevice) -> RoomsTrait:
     return device.v1_properties.rooms
 
 
+# Rooms from mock_data.HOME_DATA
+# {"id": 2362048, "name": "Example room 1"},
+# {"id": 2362044, "name": "Example room 2"},
+# {"id": 2362041, "name": "Example room 3"},
+@pytest.mark.parametrize(
+    ("room_mapping_data"),
+    [
+        ([[16, "2362048"], [17, "2362044"], [18, "2362041"]]),
+        ([[16, "2362048", 6], [17, "2362044", 14], [18, "2362041", 13]]),
+    ]
+)
 async def test_refresh_rooms_trait(
     rooms_trait: RoomsTrait,
     mock_rpc_channel: AsyncMock,
+    room_mapping_data: list[Any],
 ) -> None:
     """Test successfully getting room mapping."""
     # Setup mock to return the sample room mapping
-    mock_rpc_channel.send_command.side_effect = [
-        ROOM_MAPPING_DATA,
-    ]
+    mock_rpc_channel.send_command.side_effect = [room_mapping_data]
     # Before refresh, rooms should be empty
     assert not rooms_trait.rooms