python-kasa · rytilahti · Nov 1, 2024 · Sep 28, 2024 · Nov 1, 2024 · Nov 1, 2024
diff --git a/kasa/module.py b/kasa/module.py
@@ -127,6 +127,9 @@ class Module(ABC):
     WaterleakSensor: Final[ModuleName[smart.WaterleakSensor]] = ModuleName(
         "WaterleakSensor"
     )
+    ChildProtection: Final[ModuleName[smart.ChildProtection]] = ModuleName(
+        "ChildProtection"
+    )
     TriggerLogs: Final[ModuleName[smart.TriggerLogs]] = ModuleName("TriggerLogs")
 
     # SMARTCAMERA only modules

diff --git a/kasa/smart/modules/__init__.py b/kasa/smart/modules/__init__.py
@@ -6,6 +6,7 @@
 from .batterysensor import BatterySensor
 from .brightness import Brightness
 from .childdevice import ChildDevice
+from .childprotection import ChildProtection
 from .cloud import Cloud
 from .color import Color
 from .colortemperature import ColorTemperature
@@ -40,6 +41,7 @@
     "HumiditySensor",
     "TemperatureSensor",
     "TemperatureControl",
+    "ChildProtection",
     "ReportMode",
     "AutoOff",
     "Led",

diff --git a/kasa/smart/modules/childprotection.py b/kasa/smart/modules/childprotection.py
@@ -0,0 +1,41 @@
+"""Child lock module."""
+
+from __future__ import annotations
+
+from ...feature import Feature
+from ..smartmodule import SmartModule
+
+
+class ChildProtection(SmartModule):
+    """Implementation for child_protection."""
+
+    REQUIRED_COMPONENT = "child_protection"
+    QUERY_GETTER_NAME = "get_child_protection"
+
+    def _initialize_features(self):
+        """Initialize features after the initial update."""
+        self._add_feature(
+            Feature(
+                device=self._device,
+                id="child_lock",
+                name="Child lock",
+                container=self,
+                attribute_getter="enabled",
+                attribute_setter="set_enabled",
+                type=Feature.Type.Switch,
+                category=Feature.Category.Config,
+            )
+        )
+
+    def query(self) -> dict:
+        """Query to execute during the update cycle."""
+        return {}
+
+    @property
+    def enabled(self) -> bool:
+        """Return True if child protection is enabled."""
+        return self.data["child_protection"]
+
+    async def set_enabled(self, enabled: bool) -> dict:
+        """Set child protection."""
+        return await self.call("set_child_protection", {"enable": enabled})
diff --git a/kasa/tests/fakeprotocol_smart.py b/kasa/tests/fakeprotocol_smart.py
@@ -430,14 +430,24 @@ def _edit_preset_rules(self, info, params):
         info["get_preset_rules"]["states"][params["index"]] = params["state"]
         return {"error_code": 0}
 
+    def _update_sysinfo_key(self, info: dict, key: str, value: str) -> dict:
+        """Update a single key in the main system info.
+
+        This is used to implement child device setters that change the main sysinfo state.
+        """
+        sys_info = info.get("get_device_info", info)
+        sys_info[key] = value
+
+        return {"error_code": 0}
+
     async def _send_request(self, request_dict: dict):
         method = request_dict["method"]
 
         info = self.info
         if method == "control_child":
             return await self._handle_control_child(request_dict["params"])
 
-        params = request_dict.get("params")
+        params = request_dict.get("params", {})
         if method == "component_nego" or method[:4] == "get_":
             if method in info:
                 result = copy.deepcopy(info[method])
@@ -518,6 +528,8 @@ async def _send_request(self, request_dict: dict):
             return self._edit_preset_rules(info, params)
         elif method == "set_on_off_gradually_info":
             return self._set_on_off_gradually_info(info, params)
+        elif method == "set_child_protection":
+            return self._update_sysinfo_key(info, "child_protection", params["enable"])
         elif method[:4] == "set_":
             target_method = f"get_{method[4:]}"
             info[target_method].update(params)

diff --git a/kasa/tests/smart/modules/test_childprotection.py b/kasa/tests/smart/modules/test_childprotection.py
@@ -0,0 +1,43 @@
+import pytest
+
+from kasa import Module
+from kasa.smart.modules import ChildProtection
+from kasa.tests.device_fixtures import parametrize
+
+child_protection = parametrize(
+    "has child protection",
+    component_filter="child_protection",
+    protocol_filter={"SMART.CHILD"},
+)
+
+
+@child_protection
+@pytest.mark.parametrize(
+    ("feature", "prop_name", "type"),
+    [
+        ("child_lock", "enabled", bool),
+    ],
+)
+async def test_features(dev, feature, prop_name, type):
+    """Test that features are registered and work as expected."""
+    protect: ChildProtection = dev.modules[Module.ChildProtection]
+    assert protect is not None
+
+    prop = getattr(protect, prop_name)
+    assert isinstance(prop, type)
+
+    feat = protect._device.features[feature]
+    assert feat.value == prop
+    assert isinstance(feat.value, type)
+
+
+@child_protection
+async def test_enabled(dev):
+    """Test the API."""
+    protect: ChildProtection = dev.modules[Module.ChildProtection]
+    assert protect is not None
+
+    assert isinstance(protect.enabled, bool)
+    await protect.set_enabled(False)
+    await dev.update()
+    assert protect.enabled is False