From 5599756d289f013c7617e046cb493f133cf0fb9c Mon Sep 17 00:00:00 2001
From: Teemu R <tpr@iki.fi>
Date: Tue, 30 Apr 2024 17:31:47 +0200
Subject: [PATCH] Add support for waterleak sensor (T300) (#876)

---
 kasa/smart/modules/__init__.py             |  2 +
 kasa/smart/modules/waterleak.py            | 62 ++++++++++++++++++++++
 kasa/smart/smartchilddevice.py             |  1 +
 kasa/tests/smart/modules/test_waterleak.py | 42 +++++++++++++++
 4 files changed, 107 insertions(+)
 create mode 100644 kasa/smart/modules/waterleak.py
 create mode 100644 kasa/tests/smart/modules/test_waterleak.py

diff --git a/kasa/smart/modules/__init__.py b/kasa/smart/modules/__init__.py
index b3b1d9f4..d028b9d7 100644
--- a/kasa/smart/modules/__init__.py
+++ b/kasa/smart/modules/__init__.py
@@ -19,6 +19,7 @@ from .reportmodule import ReportModule
 from .temperature import TemperatureSensor
 from .temperaturecontrol import TemperatureControl
 from .timemodule import TimeModule
+from .waterleak import WaterleakSensor
 
 __all__ = [
     "AlarmModule",
@@ -40,4 +41,5 @@ __all__ = [
     "LightTransitionModule",
     "ColorTemperatureModule",
     "ColorModule",
+    "WaterleakSensor",
 ]
diff --git a/kasa/smart/modules/waterleak.py b/kasa/smart/modules/waterleak.py
new file mode 100644
index 00000000..1809c556
--- /dev/null
+++ b/kasa/smart/modules/waterleak.py
@@ -0,0 +1,62 @@
+"""Implementation of waterleak module."""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import TYPE_CHECKING
+
+from ...feature import Feature
+from ..smartmodule import SmartModule
+
+if TYPE_CHECKING:
+    from ..smartdevice import SmartDevice
+
+
+class WaterleakStatus(Enum):
+    """Waterleawk status."""
+
+    Normal = "normal"
+    LeakDetected = "water_leak"
+    Drying = "water_dry"
+
+
+class WaterleakSensor(SmartModule):
+    """Implementation of waterleak module."""
+
+    REQUIRED_COMPONENT = "sensor_alarm"
+
+    def __init__(self, device: SmartDevice, module: str):
+        super().__init__(device, module)
+        self._add_feature(
+            Feature(
+                device,
+                "Water leak",
+                container=self,
+                attribute_getter="status",
+                icon="mdi:water",
+            )
+        )
+        self._add_feature(
+            Feature(
+                device,
+                "Water alert",
+                container=self,
+                attribute_getter="alert",
+                icon="mdi:water-alert",
+            )
+        )
+
+    def query(self) -> dict:
+        """Query to execute during the update cycle."""
+        # Water leak information is contained in the main device info response.
+        return {}
+
+    @property
+    def status(self) -> WaterleakStatus:
+        """Return current humidity in percentage."""
+        return WaterleakStatus(self._device.sys_info["water_leak_status"])
+
+    @property
+    def alert(self) -> bool:
+        """Return true if alarm is active."""
+        return self._device.sys_info["in_alarm"]
diff --git a/kasa/smart/smartchilddevice.py b/kasa/smart/smartchilddevice.py
index 8852262c..7f747b84 100644
--- a/kasa/smart/smartchilddevice.py
+++ b/kasa/smart/smartchilddevice.py
@@ -50,6 +50,7 @@ class SmartChildDevice(SmartDevice):
         child_device_map = {
             "plug.powerstrip.sub-plug": DeviceType.Plug,
             "subg.trigger.temp-hmdt-sensor": DeviceType.Sensor,
+            "subg.trigger.water-leak-sensor": DeviceType.Sensor,
             "kasa.switch.outlet.sub-fan": DeviceType.Fan,
             "kasa.switch.outlet.sub-dimmer": DeviceType.Dimmer,
             "subg.trv": DeviceType.Thermostat,
diff --git a/kasa/tests/smart/modules/test_waterleak.py b/kasa/tests/smart/modules/test_waterleak.py
new file mode 100644
index 00000000..247ffb81
--- /dev/null
+++ b/kasa/tests/smart/modules/test_waterleak.py
@@ -0,0 +1,42 @@
+from enum import Enum
+
+import pytest
+
+from kasa.smart.modules import WaterleakSensor
+from kasa.tests.device_fixtures import parametrize
+
+waterleak = parametrize(
+    "has waterleak", component_filter="sensor_alarm", protocol_filter={"SMART.CHILD"}
+)
+
+
+@waterleak
+@pytest.mark.parametrize(
+    "feature, type",
+    [
+        ("alert", int),
+        ("status", Enum),
+    ],
+)
+async def test_waterleak_properties(dev, feature, type):
+    """Test that features are registered and work as expected."""
+    waterleak: WaterleakSensor = dev.modules["WaterleakSensor"]
+
+    prop = getattr(waterleak, feature)
+    assert isinstance(prop, type)
+
+    feat = waterleak._module_features[feature]
+    assert feat.value == prop
+    assert isinstance(feat.value, type)
+
+
+@waterleak
+async def test_waterleak_features(dev):
+    """Test waterleak features."""
+    waterleak: WaterleakSensor = dev.modules["WaterleakSensor"]
+
+    assert "water_leak" in dev.features
+    assert dev.features["water_leak"].value == waterleak.status
+
+    assert "water_alert" in dev.features
+    assert dev.features["water_alert"].value == waterleak.alert