Add multi map handling to roborock

miio/integrations/vacuum/roborock/tests/test_vacuum.py

@@ -38,6 +38,8 @@ def __init__(self, *args, **kwargs):
             "msg_seq": 320,
             "water_box_status": 1,
         }
+        self._maps = None
+        self._map_enum_cache = None
 
         self.dummies = {}
         self.dummies["consumables"] = [
@@ -71,6 +73,36 @@ def __init__(self, *args, **kwargs):
                 "end_hour": 8,
             }
         ]
+        self.dummies["multi_maps"] = [
+            {
+                "max_multi_map": 4,
+                "max_bak_map": 1,
+                "multi_map_count": 3,
+                "map_info": [
+                    {
+                        "mapFlag": 0,
+                        "add_time": 1664448893,
+                        "length": 10,
+                        "name": "Downstairs",
+                        "bak_maps": [{"mapFlag": 4, "add_time": 1663577737}],
+                    },
+                    {
+                        "mapFlag": 1,
+                        "add_time": 1663580330,
+                        "length": 8,
+                        "name": "Upstairs",
+                        "bak_maps": [{"mapFlag": 5, "add_time": 1663577752}],
+                    },
+                    {
+                        "mapFlag": 2,
+                        "add_time": 1663580384,
+                        "length": 5,
+                        "name": "Attic",
+                        "bak_maps": [{"mapFlag": 6, "add_time": 1663577765}],
+                    },
+                ],
+            }
+        ]
 
         self.return_values = {
             "get_status": lambda x: [self.state],
@@ -86,6 +118,7 @@ def __init__(self, *args, **kwargs):
             "miIO.info": "dummy info",
             "get_clean_record": lambda x: [[1488347071, 1488347123, 16, 0, 0, 0]],
             "get_dnd_timer": lambda x: self.dummies["dnd_timer"],
+            "get_multi_maps_list": lambda x: self.dummies["multi_maps"],
         }
 
         super().__init__(args, kwargs)
@@ -311,6 +344,50 @@ def test_history_empty(self):
 
             assert len(self.device.clean_history().ids) == 0
 
+    def test_get_maps_dict(self):
+        MAP_LIST = [
+            {
+                "mapFlag": 0,
+                "add_time": 1664448893,
+                "length": 10,
+                "name": "Downstairs",
+                "bak_maps": [{"mapFlag": 4, "add_time": 1663577737}],
+            },
+            {
+                "mapFlag": 1,
+                "add_time": 1663580330,
+                "length": 8,
+                "name": "Upstairs",
+                "bak_maps": [{"mapFlag": 5, "add_time": 1663577752}],
+            },
+            {
+                "mapFlag": 2,
+                "add_time": 1663580384,
+                "length": 5,
+                "name": "Attic",
+                "bak_maps": [{"mapFlag": 6, "add_time": 1663577765}],
+            },
+        ]
+
+        with patch.object(
+            self.device,
+            "send",
+            return_value=[
+                {
+                    "max_multi_map": 4,
+                    "max_bak_map": 1,
+                    "multi_map_count": 3,
+                    "map_info": MAP_LIST,
+                }
+            ],
+        ):
+            maps = self.device.get_maps()
+
+        assert maps.map_count == 3
+        assert maps.map_id_list == [0, 1, 2]
+        assert maps.map_list == MAP_LIST
+        assert maps.map_name_dict == {"Downstairs": 0, "Upstairs": 1, "Attic": 2}
+
     def test_info_no_cloud(self):
         """Test the info functionality for non-cloud connected device."""
         from miio.exceptions import DeviceInfoUnavailableException

miio/integrations/vacuum/roborock/vacuum.py

@@ -1,5 +1,6 @@
 import contextlib
 import datetime
+import enum
 import json
 import logging
 import math
@@ -46,6 +47,7 @@
     CleaningSummary,
     ConsumableStatus,
     DNDStatus,
+    MapList,
     SoundInstallStatus,
     SoundStatus,
     Timer,
@@ -135,6 +137,8 @@ def __init__(
             ip, token, start_id, debug, lazy_discover, timeout, model=model
         )
         self.manual_seqnum = -1
+        self._maps: Optional[MapList] = None
+        self._map_enum_cache = None
 
     @command()
     def start(self):
@@ -365,6 +369,40 @@ def map(self):
         # returns ['retry'] without internet
         return self.send("get_map_v1")
 
+    @command()
+    def get_maps(self) -> MapList:
+        """Return list of maps."""
+        if self._maps is not None:
+            return self._maps
+
+        self._maps = MapList(self.send("get_multi_maps_list")[0])
+        return self._maps
+
+    def _map_enum(self) -> Optional[enum.Enum]:
+        """Enum of the available map names."""
+        if self._map_enum_cache is not None:
+            return self._map_enum_cache
+
+        maps = self.get_maps()
+
+        self._map_enum_cache = enum.Enum("map_enum", maps.map_name_dict)
+        return self._map_enum_cache
+
+    @command(click.argument("map_id", type=int))
+    def load_map(
+        self,
+        map_enum: Optional[enum.Enum] = None,
+        map_id: Optional[int] = None,
+    ):
+        """Change the current map used."""
+        if map_enum is None and map_id is None:
+            raise ValueError("Either map_enum or map_id is required.")
+
+        if map_enum is not None:
+            map_id = map_enum.value
+
+        return self.send("load_multi_map", [map_id])[0] == "ok"
+
     @command(click.argument("start", type=bool))
     def edit_map(self, start):
         """Start map editing?"""

miio/integrations/vacuum/roborock/vacuumcontainers.py

@@ -1,3 +1,4 @@
+import logging
 from datetime import datetime, time, timedelta
 from enum import IntEnum
 from typing import Any, Dict, List, Optional, Union
@@ -12,6 +13,8 @@
 
 from .vacuum_enums import MopIntensity, MopMode
 
+_LOGGER = logging.getLogger(__name__)
+
 
 def pretty_area(x: float) -> float:
     return int(x) / 1000000
@@ -94,6 +97,42 @@ def pretty_area(x: float) -> float:
 }
 
 
+class MapList(DeviceStatus):
+    """Contains a information about the maps/floors of the vacuum."""
+
+    def __init__(self, data: Dict[str, Any]) -> None:
+        # {'max_multi_map': 4, 'max_bak_map': 1, 'multi_map_count': 3, 'map_info': [
+        #    {'mapFlag': 0, 'add_time': 1664448893, 'length': 10, 'name': 'Downstairs', 'bak_maps': [{'mapFlag': 4, 'add_time': 1663577737}]},
+        #    {'mapFlag': 1, 'add_time': 1663580330, 'length': 8, 'name': 'Upstairs', 'bak_maps': [{'mapFlag': 5, 'add_time': 1663577752}]},
+        #    {'mapFlag': 2, 'add_time': 1663580384, 'length': 5, 'name': 'Attic', 'bak_maps': [{'mapFlag': 6, 'add_time': 1663577765}]}
+        #  ]}
+        self.data = data
+
+        self._map_name_dict = {}
+        for map in self.data["map_info"]:
+            self._map_name_dict[map["name"]] = map["mapFlag"]
+
+    @property
+    def map_count(self) -> int:
+        """Amount of maps stored."""
+        return self.data["multi_map_count"]
+
+    @property
+    def map_id_list(self) -> List[int]:
+        """List of map ids."""
+        return list(self._map_name_dict.values())
+
+    @property
+    def map_list(self) -> List[Dict[str, Any]]:
+        """List of map info."""
+        return self.data["map_info"]
+
+    @property
+    def map_name_dict(self) -> Dict[str, int]:
+        """Dictionary of map names (keys) with there ids (values)."""
+        return self._map_name_dict
+
+
 class VacuumStatus(VacuumDeviceStatus):
     """Container for status reports from the vacuum."""
 
@@ -284,6 +323,20 @@ def map(self) -> bool:
         """Map token."""
         return bool(self.data["map_present"])
 
+    @property
+    @setting(
+        "Current map",
+        choices_attribute="_map_enum",
+        setter_name="load_map",
+        icon="mdi:floor-plan",
+    )
+    def current_map_id(self) -> int:
+        """The id of the current map with regards to the multi map feature,
+
+        [3,7,11,15] -> [0,1,2,3].
+        """
+        return int((self.data["map_status"] + 1) / 4 - 1)
+
     @property
     def in_zone_cleaning(self) -> bool:
         """Return True if the vacuum is in zone cleaning mode."""
@@ -502,6 +555,11 @@ def area(self) -> float:
         """Total cleaned area."""
         return pretty_area(self.data["area"])
 
+    @property
+    def map_id(self) -> int:
+        """Map id used (multi map feature) during the cleaning run."""
+        return self.data.get("map_flag", 0)
+
     @property
     def error_code(self) -> int:
         """Error code."""