Add setting to change clean count (#1457)

rytilahti · web-flow · commit bc97c0794a1b · 2025-01-15T19:11:33.000+01:00
Adds a setting to change the number of times to clean:
```
== Configuration ==
Clean count (clean_count): 1 (range: 1-3)
```
diff --git a/kasa/smart/modules/clean.py b/kasa/smart/modules/clean.py
@@ -5,7 +5,7 @@
 import logging
 from datetime import timedelta
 from enum import IntEnum
-from typing import Annotated
+from typing import Annotated, Literal
 
 from ...feature import Feature
 from ...module import FeatureAttribute
@@ -157,6 +157,19 @@ def _initialize_features(self) -> None:
                 type=Feature.Type.Choice,
             )
         )
+        self._add_feature(
+            Feature(
+                self._device,
+                id="clean_count",
+                name="Clean count",
+                container=self,
+                attribute_getter="clean_count",
+                attribute_setter="set_clean_count",
+                range_getter=lambda: (1, 3),
+                category=Feature.Category.Config,
+                type=Feature.Type.Number,
+            )
+        )
         self._add_feature(
             Feature(
                 self._device,
@@ -283,9 +296,17 @@ async def set_fan_speed_preset(
         name_to_value = {x.name: x.value for x in FanSpeed}
         if speed not in name_to_value:
             raise ValueError("Invalid fan speed %s, available %s", speed, name_to_value)
-        return await self.call(
-            "setCleanAttr", {"suction": name_to_value[speed], "type": "global"}
-        )
+        return await self._change_setting("suction", name_to_value[speed])
+
+    async def _change_setting(
+        self, name: str, value: int, *, scope: Literal["global", "pose"] = "global"
+    ) -> dict:
+        """Change device setting."""
+        params = {
+            name: value,
+            "type": scope,
+        }
+        return await self.call("setCleanAttr", params)
 
     @property
     def battery(self) -> int:
@@ -339,3 +360,12 @@ def clean_time(self) -> timedelta:
     def clean_progress(self) -> int:
         """Return amount of currently cleaned area."""
         return self._info["clean_percent"]
+
+    @property
+    def clean_count(self) -> Annotated[int, FeatureAttribute()]:
+        """Return number of times to clean."""
+        return self._settings["clean_number"]
+
+    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)
diff --git a/tests/smart/modules/test_clean.py b/tests/smart/modules/test_clean.py
@@ -69,6 +69,13 @@ async def test_features(dev: SmartDevice, feature: str, prop_name: str, type: ty
             {"suction": 1, "type": "global"},
             id="vacuum_fan_speed",
         ),
+        pytest.param(
+            "clean_count",
+            2,
+            "setCleanAttr",
+            {"clean_number": 2, "type": "global"},
+            id="clean_count",
+        ),
     ],
 )
 @clean
