From 06f14bfc6d715ea0db68850310164c17d57a9af8 Mon Sep 17 00:00:00 2001
From: Teemu Rytilahti <tpr@iki.fi>
Date: Wed, 9 Oct 2024 16:00:50 +0200
Subject: [PATCH] Add alert volume setting to water leak sensor

---
 kasa/cli/main.py                           |  2 +-
 kasa/smart/modules/waterleaksensor.py      | 32 ++++++++++++++++++----
 kasa/tests/smart/modules/test_waterleak.py |  1 +
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/kasa/cli/main.py b/kasa/cli/main.py
index 88b768c4..c0e66f40 100755
--- a/kasa/cli/main.py
+++ b/kasa/cli/main.py
@@ -339,7 +339,7 @@ async def cli(
     # Skip update on specific commands, or if device factory,
     # that performs an update was used for the device.
     if ctx.invoked_subcommand not in SKIP_UPDATE_COMMANDS and not device_updated:
-        await dev.update()
+        await dev.update(update_children=True)
 
     @asynccontextmanager
     async def async_wrapped_device(device: Device):
diff --git a/kasa/smart/modules/waterleaksensor.py b/kasa/smart/modules/waterleaksensor.py
index bba4f61d..3c5d8b0f 100644
--- a/kasa/smart/modules/waterleaksensor.py
+++ b/kasa/smart/modules/waterleaksensor.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from enum import Enum
+from typing import Literal, TypeAlias
 
 from ...feature import Feature
 from ..smartmodule import SmartModule
@@ -16,10 +17,15 @@ class WaterleakStatus(Enum):
     Drying = "water_dry"
 
 
+Volume: TypeAlias = Literal["low", "normal", "high", "mute"]
+ALLOWED_VOLUMES = ["low", "normal", "high", "mute"]
+
+
 class WaterleakSensor(SmartModule):
     """Implementation of waterleak module."""
 
     REQUIRED_COMPONENT = "sensor_alarm"
+    QUERY_GETTER_NAME = "get_alarm_config"
 
     def _initialize_features(self):
         """Initialize features after the initial update."""
@@ -47,11 +53,18 @@ class WaterleakSensor(SmartModule):
                 type=Feature.Type.BinarySensor,
             )
         )
-
-    def query(self) -> dict:
-        """Query to execute during the update cycle."""
-        # Water leak information is contained in the main device info response.
-        return {}
+        self._add_feature(
+            Feature(
+                self._device,
+                id="water_alert_volume",
+                name="Water alert volume",
+                container=self,
+                attribute_getter="alert_volume",
+                attribute_setter="set_alert_volume",
+                type=Feature.Type.Choice,
+                choices_getter=lambda: ALLOWED_VOLUMES,
+            )
+        )
 
     @property
     def status(self) -> WaterleakStatus:
@@ -62,3 +75,12 @@ class WaterleakSensor(SmartModule):
     def alert(self) -> bool:
         """Return true if alarm is active."""
         return self._device.sys_info["in_alarm"]
+
+    @property
+    def alert_volume(self) -> Volume:
+        """Get water leak alert volume."""
+        return self.data["volume"]
+
+    async def set_alert_volume(self, volume: Volume):
+        """Set water leak alert volume."""
+        await self.call("set_alarm_config", {"volume": volume})
diff --git a/kasa/tests/smart/modules/test_waterleak.py b/kasa/tests/smart/modules/test_waterleak.py
index c48d8244..679b2923 100644
--- a/kasa/tests/smart/modules/test_waterleak.py
+++ b/kasa/tests/smart/modules/test_waterleak.py
@@ -16,6 +16,7 @@ waterleak = parametrize(
     [
         ("water_alert", "alert", int),
         ("water_leak", "status", Enum),
+        ("water_alert_volume", "alert_volume", str),
     ],
 )
 async def test_waterleak_properties(dev, feature, prop_name, type):