python-kasa · rytilahti · Jan 15, 2025 · Jan 15, 2025 · Jan 15, 2025 · Jan 15, 2025
diff --git a/kasa/module.py b/kasa/module.py
@@ -165,6 +165,7 @@ class Module(ABC):
     Clean: Final[ModuleName[smart.Clean]] = ModuleName("Clean")
     Dustbin: Final[ModuleName[smart.Dustbin]] = ModuleName("Dustbin")
     Speaker: Final[ModuleName[smart.Speaker]] = ModuleName("Speaker")
+    Mop: Final[ModuleName[smart.Mop]] = ModuleName("Mop")
 
     def __init__(self, device: Device, module: str) -> None:
         self._device = device

diff --git a/kasa/smart/modules/__init__.py b/kasa/smart/modules/__init__.py
@@ -27,6 +27,7 @@
 from .lightstripeffect import LightStripEffect
 from .lighttransition import LightTransition
 from .matter import Matter
+from .mop import Mop
 from .motionsensor import MotionSensor
 from .overheatprotection import OverheatProtection
 from .reportmode import ReportMode
@@ -76,4 +77,5 @@
     "HomeKit",
     "Matter",
     "Dustbin",
+    "Mop",
 ]
diff --git a/kasa/smart/modules/mop.py b/kasa/smart/modules/mop.py
@@ -0,0 +1,90 @@
+"""Implementation of vacuum mop."""
+
+from __future__ import annotations
+
+import logging
+from enum import IntEnum
+from typing import Annotated
+
+from ...feature import Feature
+from ...module import FeatureAttribute
+from ..smartmodule import SmartModule
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class Waterlevel(IntEnum):
+    """Water level for mopping."""
+
+    Disable = 0
+    Low = 1
+    Medium = 2
+    High = 3
+
+
+class Mop(SmartModule):
+    """Implementation of vacuum mop."""
+
+    REQUIRED_COMPONENT = "mop"
+
+    def _initialize_features(self) -> None:
+        """Initialize features."""
+        self._add_feature(
+            Feature(
+                self._device,
+                id="mop_attached",
+                name="Mop attached",
+                container=self,
+                icon="mdi:square-rounded",
+                attribute_getter="mop_attached",
+                category=Feature.Category.Info,
+                type=Feature.BinarySensor,
+            )
+        )
+
+        self._add_feature(
+            Feature(
+                self._device,
+                id="mop_waterlevel",
+                name="Mop water level",
+                container=self,
+                attribute_getter="waterlevel",
+                attribute_setter="set_waterlevel",
+                icon="mdi:water",
+                choices_getter=lambda: list(Waterlevel.__members__),
+                category=Feature.Category.Config,
+                type=Feature.Type.Choice,
+            )
+        )
+
+    def query(self) -> dict:
+        """Query to execute during the update cycle."""
+        return {
+            "getMopState": {},
+            "getCleanAttr": {"type": "global"},
+        }
+
+    @property
+    def mop_attached(self) -> bool:
+        """Return True if mop is attached."""
+        return self.data["getMopState"]["mop_state"]
+
+    @property
+    def _settings(self) -> dict:
+        """Return settings settings."""
+        return self.data["getCleanAttr"]
+
+    @property
+    def waterlevel(self) -> Annotated[str, FeatureAttribute()]:
+        """Return water level."""
+        return Waterlevel(int(self._settings["cistern"])).name
+
+    async def set_waterlevel(self, mode: str) -> Annotated[dict, FeatureAttribute()]:
+        """Set waterlevel mode."""
+        name_to_value = {x.name: x.value for x in Waterlevel}
+        if mode not in name_to_value:
+            raise ValueError("Invalid waterlevel %s, available %s", mode, name_to_value)
+
+        settings = self._settings.copy()
+        settings["cistern"] = name_to_value[mode]
+        return await self.call("setCleanAttr", settings)
diff --git a/tests/smart/modules/test_mop.py b/tests/smart/modules/test_mop.py
@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import pytest
+from pytest_mock import MockerFixture
+
+from kasa import Module
+from kasa.smart import SmartDevice
+from kasa.smart.modules.mop import Waterlevel
+
+from ...device_fixtures import get_parent_and_child_modules, parametrize
+
+mop = parametrize("has mop", component_filter="mop", protocol_filter={"SMART"})
+
+
+@mop
+@pytest.mark.parametrize(
+    ("feature", "prop_name", "type"),
+    [
+        ("mop_attached", "mop_attached", bool),
+        ("mop_waterlevel", "waterlevel", str),
+    ],
+)
+async def test_features(dev: SmartDevice, feature: str, prop_name: str, type: type):
+    """Test that features are registered and work as expected."""
+    mod = next(get_parent_and_child_modules(dev, Module.Mop))
+    assert mod is not None
+
+    prop = getattr(mod, prop_name)
+    assert isinstance(prop, type)
+
+    feat = mod._device.features[feature]
+    assert feat.value == prop
+    assert isinstance(feat.value, type)
+
+
+@mop
+async def test_mop_waterlevel(dev: SmartDevice, mocker: MockerFixture):
+    """Test dust mode."""
+    mop_module = next(get_parent_and_child_modules(dev, Module.Mop))
+    call = mocker.spy(mop_module, "call")
+
+    waterlevel = mop_module._device.features["mop_waterlevel"]
+    assert mop_module.waterlevel == waterlevel.value
+
+    new_level = Waterlevel.High
+    await mop_module.set_waterlevel(new_level.name)
+
+    params = mop_module._settings.copy()
+    params["cistern"] = new_level.value
+
+    call.assert_called_with("setCleanAttr", params)
+
+    await dev.update()
+
+    assert mop_module.waterlevel == new_level.name
+
+    with pytest.raises(ValueError, match="Invalid waterlevel"):
+        await mop_module.set_waterlevel("invalid")