Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
144 changes: 142 additions & 2 deletions kasa/smart/modules/clean.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -13,6 +15,10 @@

_LOGGER = logging.getLogger(__name__)

# Only known value for start_type in setSwitchClean; required for
# targeted cleaning modes (Room) but not for StandardHome.
_START_TYPE_RESUME = 1


class Status(IntEnum):
"""Status of vacuum."""
Expand Down Expand Up @@ -58,6 +64,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 (matches :class:`FanSpeed` values).
suction: int = 0
Comment on lines +91 to +92
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the comment is correct, maybe the type of this should be FanSpeed (assuming it's an intenum, currently on mobile)?

#: 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."""

Expand Down Expand Up @@ -263,6 +328,7 @@ def query(self) -> dict:
"getBatteryInfo": {},
"getCleanStatus": {},
"getCleanAttr": {"type": "global"},
"getMapInfo": {},
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would you mind adding this query also to the dump_devinfos query list, and creating a fixture with the data? That will make it testable and easier to extend in the future.

}

async def start(self) -> dict:
Expand All @@ -275,7 +341,7 @@ async def start(self) -> dict:
return await self.call(
"setSwitchClean",
{
"clean_mode": 0,
"clean_mode": CleanMode.StandardHome,
"clean_on": True,
"clean_order": True,
"force_clean": False,
Expand Down Expand Up @@ -409,3 +475,77 @@ def clean_count(self) -> Annotated[int, FeatureAttribute()]:
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:
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we also have the name information? If yes, maybe it makes sense to expose it also?

"""Return the ID of the currently active map."""
return self.data["getMapInfo"]["current_map_id"]

@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_RESUME,
},
)

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": 0})
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any info about the type and its values? Enumization would be nice here, too.


rooms: list[RoomInfo] = []
for area in resp.get("area_list", []):
if area.get("type") != AreaType.Room:
continue
Comment on lines +533 to +534
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Adding debug logging here could be useful when extending this.

name = None
if raw_name := area.get("name"):
try:
name = base64.b64decode(raw_name).decode()
except Exception:
name = raw_name
rooms.append(
Comment on lines +540 to +541
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add a newline in-between to create a logical separation.

RoomInfo(
id=area["id"],
name=name,
color=area.get("color", 0),
suction=area.get("suction", 0),
cistern=area.get("cistern", 0),
clean_number=area.get("clean_number", 0),
)
)
return rooms
179 changes: 178 additions & 1 deletion tests/smart/modules/test_clean.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

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, RoomInfo, Status

from ...device_fixtures import get_parent_and_child_modules, parametrize

Expand Down Expand Up @@ -239,3 +239,180 @@ async def test_invalid_settings(

with pytest.raises(exc, match=exc_message):
await setter(value)


@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[1].id == 3
assert rooms[1].name == "Living Room"


@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


@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