Add selective room cleaning to vacuum clean module (#1660)

Adds support for selective room cleaning to the existing `Clean` module.

Co-authored-by: davidefiocco <davidefiocco@users.noreply.github.com>

kasa/smart/modules/clean.py

@@ -2,9 +2,11 @@
 
 from __future__ import annotations
 
+import base64
 import logging
+from dataclasses import dataclass
 from datetime import timedelta
-from enum import IntEnum
+from enum import IntEnum, StrEnum
 from typing import Annotated, Literal
 
 from ...feature import Feature
@@ -13,6 +15,13 @@ from ..smartmodule import SmartModule
 
 _LOGGER = logging.getLogger(__name__)
 
+# Only known value for start_type in setSwitchClean; required for
+# targeted cleaning modes (Room) but not for StandardHome.
+_START_TYPE_DEFAULT = 1
+
+# Only known value for the type parameter in getMapData.
+_MAP_DATA_TYPE_DEFAULT = 0
+
 
 class Status(IntEnum):
     """Status of vacuum."""
@@ -58,6 +67,65 @@ class FanSpeed(IntEnum):
     Ultra = 5
 
 
+class CleanMode(IntEnum):
+    """Clean mode for ``setSwitchClean`` and ``getCleanStatus``.
+
+    Used as ``clean_mode`` in commands and ``clean_status`` in status responses.
+    """
+
+    #: Clean all rooms with uniform settings.
+    StandardHome = 0
+    #: Clean all rooms with per-room settings and custom order.
+    AdvancedHome = 1
+    #: Clean a small area around the vacuum's current position.
+    Spot = 2
+    #: Clean selected rooms only.
+    Room = 3
+    #: Clean user-defined rectangular areas.
+    Zone = 4
+    #: Run a saved custom cleaning preset.
+    Custom = 5
+
+
+@dataclass
+class CleanAreaSettings:
+    """Per-area cleaning settings shared by rooms and zones."""
+
+    #: Suction power level, or ``None`` when not configured per-area.
+    suction: FanSpeed | None = None
+    #: Water level for mopping.
+    cistern: int = 0
+    #: Number of cleaning passes.
+    clean_number: int = 0
+
+
+@dataclass
+class RoomInfo(CleanAreaSettings):
+    """Information about a room on the vacuum's map."""
+
+    #: Room ID used in cleaning commands.
+    id: int = 0
+    #: Human-readable room name (base64-decoded from the device).
+    name: str | None = None
+    #: Color index used for map rendering.
+    color: int = 0
+
+
+class AreaType(StrEnum):
+    """Type of area entry in map data."""
+
+    #: A named room.
+    Room = "room"
+    #: A user-defined rectangular cleaning zone.
+    Area = "area"
+    #: A virtual wall boundary.
+    VirtualWall = "virtual_wall"
+    #: A no-go zone.
+    Forbid = "forbid"
+    #: A detected carpet region.
+    CarpetRectangle = "carpet_rectangle"
+
+
 class AreaUnit(IntEnum):
     """Area unit."""
 
@@ -263,6 +331,7 @@ class Clean(SmartModule):
             "getBatteryInfo": {},
             "getCleanStatus": {},
             "getCleanAttr": {"type": "global"},
+            "getMapInfo": {},
         }
 
     async def start(self) -> dict:
@@ -275,7 +344,7 @@ class Clean(SmartModule):
         return await self.call(
             "setSwitchClean",
             {
-                "clean_mode": 0,
+                "clean_mode": CleanMode.StandardHome,
                 "clean_on": True,
                 "clean_order": True,
                 "force_clean": False,
@@ -409,3 +478,102 @@ class Clean(SmartModule):
     async def set_clean_count(self, count: int) -> Annotated[dict, FeatureAttribute()]:
         """Set number of times to clean."""
         return await self._change_setting("clean_number", count)
+
+    @property
+    def current_map_id(self) -> int:
+        """Return the ID of the currently active map."""
+        return self.data["getMapInfo"]["current_map_id"]
+
+    @property
+    def current_map_name(self) -> str | None:
+        """Return the name of the currently active map, or ``None``."""
+        map_info = self.data["getMapInfo"]
+        current_id = map_info["current_map_id"]
+        for m in map_info.get("map_list", []):
+            if m.get("map_id") == current_id:
+                if raw_name := m.get("map_name"):
+                    try:
+                        return base64.b64decode(raw_name).decode()
+                    except Exception:
+                        return raw_name
+                return None
+        return None
+
+    @property
+    def clean_type(self) -> CleanMode | None:
+        """Return the active cleaning mode, or ``None`` if unavailable."""
+        cs = self.data.get("getCleanStatus")
+        if cs is None or "clean_status" not in cs:
+            return None
+        return CleanMode(cs["clean_status"])
+
+    async def clean_rooms(
+        self, room_ids: list[int], *, map_id: int | None = None
+    ) -> dict:
+        """Start cleaning specific rooms.
+
+        Per-room settings are not supported; the device uses the global
+        suction / cistern / clean_number values for room cleaning.
+
+        :param room_ids: List of room IDs to clean.
+        :param map_id: Map ID to clean on. Defaults to the current active map.
+        """
+        if not room_ids:
+            raise ValueError("room_ids must not be empty")
+        if map_id is None:
+            map_id = self.current_map_id
+        return await self.call(
+            "setSwitchClean",
+            {
+                "clean_mode": CleanMode.Room,
+                "clean_on": True,
+                "clean_order": True,
+                "force_clean": False,
+                "map_id": map_id,
+                "room_list": list(room_ids),
+                "start_type": _START_TYPE_DEFAULT,
+            },
+        )
+
+    async def get_rooms(self, map_id: int | None = None) -> list[RoomInfo]:
+        """Return the list of rooms for the given map.
+
+        Room names are base64-decoded when present.
+
+        :param map_id: Map ID to query. Defaults to the current active map.
+        """
+        if map_id is None:
+            map_id = self.current_map_id
+        resp = await self.call(
+            "getMapData", {"map_id": map_id, "type": _MAP_DATA_TYPE_DEFAULT}
+        )
+        map_data = resp.get("getMapData", resp)
+
+        rooms: list[RoomInfo] = []
+        for area in map_data.get("area_list", []):
+            if area.get("type") != AreaType.Room:
+                _LOGGER.debug(
+                    "Skipping non-room area: type=%s, id=%s",
+                    area.get("type"),
+                    area.get("id"),
+                )
+                continue
+            name = None
+            if raw_name := area.get("name"):
+                try:
+                    name = base64.b64decode(raw_name).decode()
+                except Exception:
+                    name = raw_name
+
+            suction_val = area.get("suction")
+            rooms.append(
+                RoomInfo(
+                    id=area["id"],
+                    name=name,
+                    color=area.get("color", 0),
+                    suction=FanSpeed(suction_val) if suction_val else None,
+                    cistern=area.get("cistern", 0),
+                    clean_number=area.get("clean_number", 0),
+                )
+            )
+        return rooms

tests/smart/modules/test_clean.py

@@ -7,7 +7,7 @@ from pytest_mock import MockerFixture
 
 from kasa import Module
 from kasa.smart import SmartDevice
-from kasa.smart.modules.clean import ErrorCode, Status
+from kasa.smart.modules.clean import CleanMode, ErrorCode, FanSpeed, RoomInfo, Status
 
 from ...device_fixtures import get_parent_and_child_modules, parametrize
 
@@ -241,3 +241,221 @@ async def test_invalid_settings(
 
     with pytest.raises(exc, match=exc_message):
         await setter(value)
+
+
+@clean
+async def test_current_map_name(dev: SmartDevice, mocker: MockerFixture):
+    """Test current_map_name decodes the base64 name of the active map."""
+    import base64
+
+    clean = next(get_parent_and_child_modules(dev, Module.Clean))
+
+    map_info = {
+        "current_map_id": 42,
+        "map_list": [
+            {"map_id": 42, "map_name": base64.b64encode(b"Upstairs").decode()},
+            {"map_id": 99, "map_name": base64.b64encode(b"Downstairs").decode()},
+        ],
+    }
+    mocker.patch.object(
+        type(clean),
+        "data",
+        new_callable=mocker.PropertyMock,
+        return_value={**clean.data, "getMapInfo": map_info},
+    )
+    assert clean.current_map_name == "Upstairs"
+
+
+@clean
+async def test_current_map_name_missing(dev: SmartDevice, mocker: MockerFixture):
+    """Test current_map_name returns None when map_list is empty."""
+    clean = next(get_parent_and_child_modules(dev, Module.Clean))
+
+    map_info = {"current_map_id": 0, "map_list": []}
+    mocker.patch.object(
+        type(clean),
+        "data",
+        new_callable=mocker.PropertyMock,
+        return_value={**clean.data, "getMapInfo": map_info},
+    )
+    assert clean.current_map_name is None
+
+
+@clean
+async def test_clean_rooms(dev: SmartDevice, mocker: MockerFixture):
+    """Test clean_rooms sends the correct setSwitchClean payload."""
+    clean = next(get_parent_and_child_modules(dev, Module.Clean))
+    call = mocker.spy(clean, "call")
+
+    room_ids = [2, 3]
+    await clean.clean_rooms(room_ids)
+
+    call.assert_called_with(
+        "setSwitchClean",
+        {
+            "clean_mode": CleanMode.Room,
+            "clean_on": True,
+            "clean_order": True,
+            "force_clean": False,
+            "map_id": clean.current_map_id,
+            "room_list": room_ids,
+            "start_type": 1,
+        },
+    )
+
+
+@clean
+async def test_clean_rooms_explicit_map_id(dev: SmartDevice, mocker: MockerFixture):
+    """Test clean_rooms uses the provided map_id when given."""
+    clean = next(get_parent_and_child_modules(dev, Module.Clean))
+    call = mocker.spy(clean, "call")
+
+    await clean.clean_rooms([5], map_id=12345)
+
+    call.assert_called_with(
+        "setSwitchClean",
+        {
+            "clean_mode": CleanMode.Room,
+            "clean_on": True,
+            "clean_order": True,
+            "force_clean": False,
+            "map_id": 12345,
+            "room_list": [5],
+            "start_type": 1,
+        },
+    )
+
+
+@clean
+async def test_clean_rooms_empty_raises(dev: SmartDevice):
+    """Test clean_rooms raises ValueError when room_ids is empty."""
+    clean = next(get_parent_and_child_modules(dev, Module.Clean))
+
+    with pytest.raises(ValueError, match="room_ids must not be empty"):
+        await clean.clean_rooms([])
+
+
+@clean
+async def test_get_rooms(dev: SmartDevice, mocker: MockerFixture):
+    """Test get_rooms returns RoomInfo objects with decoded names."""
+    import base64
+
+    clean = next(get_parent_and_child_modules(dev, Module.Clean))
+
+    map_data = {
+        "area_list": [
+            {
+                "id": 2,
+                "name": base64.b64encode(b"Kitchen").decode(),
+                "type": "room",
+                "color": 1,
+                "suction": 2,
+                "cistern": 1,
+                "clean_number": 1,
+            },
+            {
+                "id": 3,
+                "name": base64.b64encode(b"Living Room").decode(),
+                "type": "room",
+                "color": 2,
+                "suction": 3,
+                "cistern": 2,
+                "clean_number": 2,
+            },
+            {"id": 401, "type": "virtual_wall", "vertexs": []},
+        ]
+    }
+    call_mock = mocker.patch.object(clean, "call", return_value=map_data)
+
+    rooms = await clean.get_rooms()
+
+    call_mock.assert_called_once_with(
+        "getMapData", {"map_id": clean.current_map_id, "type": 0}
+    )
+    assert len(rooms) == 2
+    assert all(isinstance(r, RoomInfo) for r in rooms)
+    assert rooms[0].id == 2
+    assert rooms[0].name == "Kitchen"
+    assert rooms[0].color == 1
+    assert rooms[0].suction is FanSpeed.Standard
+    assert rooms[1].id == 3
+    assert rooms[1].name == "Living Room"
+    assert rooms[1].suction is FanSpeed.Turbo
+
+
+@clean
+async def test_get_rooms_explicit_map_id(dev: SmartDevice, mocker: MockerFixture):
+    """Test get_rooms uses the provided map_id when given."""
+    import base64
+
+    clean = next(get_parent_and_child_modules(dev, Module.Clean))
+
+    map_data = {
+        "area_list": [
+            {
+                "id": 1,
+                "name": base64.b64encode(b"Hall").decode(),
+                "type": "room",
+            },
+        ]
+    }
+    call_mock = mocker.patch.object(clean, "call", return_value=map_data)
+
+    rooms = await clean.get_rooms(map_id=99999)
+
+    call_mock.assert_called_once_with("getMapData", {"map_id": 99999, "type": 0})
+    assert len(rooms) == 1
+    assert rooms[0].name == "Hall"
+
+
+@clean
+async def test_get_rooms_no_name(dev: SmartDevice, mocker: MockerFixture):
+    """Test get_rooms handles rooms without names."""
+    clean = next(get_parent_and_child_modules(dev, Module.Clean))
+
+    map_data = {"area_list": [{"id": 5, "type": "room"}]}
+    mocker.patch.object(clean, "call", return_value=map_data)
+
+    rooms = await clean.get_rooms()
+
+    assert len(rooms) == 1
+    assert rooms[0].id == 5
+    assert rooms[0].name is None
+    assert rooms[0].suction is None
+
+
+@clean
+async def test_clean_type(dev: SmartDevice, mocker: MockerFixture):
+    """Test clean_type returns the correct CleanMode."""
+    clean = next(get_parent_and_child_modules(dev, Module.Clean))
+
+    mocker.patch.object(
+        type(clean),
+        "data",
+        new_callable=mocker.PropertyMock,
+        return_value={**clean.data, "getCleanStatus": {"clean_status": 0}},
+    )
+    assert clean.clean_type is CleanMode.StandardHome
+
+    mocker.patch.object(
+        type(clean),
+        "data",
+        new_callable=mocker.PropertyMock,
+        return_value={**clean.data, "getCleanStatus": {"clean_status": 3}},
+    )
+    assert clean.clean_type is CleanMode.Room
+
+
+@clean
+async def test_clean_type_missing(dev: SmartDevice, mocker: MockerFixture):
+    """Test clean_type returns None when clean_status is unavailable."""
+    clean = next(get_parent_and_child_modules(dev, Module.Clean))
+
+    data_without = {k: v for k, v in clean.data.items() if k != "getCleanStatus"}
+    mocker.patch.object(
+        type(clean),
+        "data",
+        new_callable=mocker.PropertyMock,
+        return_value=data_without,
+    )
+    assert clean.clean_type is None