python-kasa · sdb9696 · Apr 30, 2024 · Apr 28, 2024 · Apr 29, 2024 · Apr 29, 2024
diff --git a/kasa/feature.py b/kasa/feature.py
@@ -92,6 +92,13 @@ class Category(Enum):
     #: If set, this property will be used to set *minimum_value* and *maximum_value*.
     range_getter: str | None = None
 
+    # Choice-specific attributes
+    #: List of choices as enum
+    choices: list[str] | None = None
+    #: Attribute name of the choices getter property.
+    #: If set, this property will be used to set *choices*.
+    choices_getter: str | None = None
+
     #: Identifier
     id: str | None = None
 
@@ -108,6 +115,10 @@ def __post_init__(self):
                 container, self.range_getter
             )
 
+        # Populate choices, if choices_getter is given
+        if self.choices_getter is not None:
+            self.choices = getattr(container, self.choices_getter)
+
         # Set the category, if unset
         if self.category is Feature.Category.Unset:
             if self.attribute_setter:
@@ -147,6 +158,12 @@ async def set_value(self, value):
                     f"Value {value} out of range "
                     f"[{self.minimum_value}, {self.maximum_value}]"
                 )
+        elif self.type == Feature.Type.Choice:  # noqa: SIM102
+            if value not in self.choices:
+                raise ValueError(
+                    f"Unexpected value for {self.name}: {value}"
+                    f" - allowed: {self.choices}"
+                )
 
         container = self.container if self.container is not None else self.device
         if self.type == Feature.Type.Action:

diff --git a/kasa/module.py b/kasa/module.py
@@ -40,6 +40,12 @@ def query(self):
     def data(self):
         """Return the module specific raw data from the last update."""
 
+    def _initialize_features(self):  # noqa: B027
+        """Initialize features after the initial update.
+
+        This can be implemented if features depend on module query responses.
+        """
+
     def _add_feature(self, feature: Feature):
         """Add module feature."""
 

diff --git a/kasa/smart/modules/alarmmodule.py b/kasa/smart/modules/alarmmodule.py
@@ -2,14 +2,9 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from ...feature import Feature
 from ..smartmodule import SmartModule
 
-if TYPE_CHECKING:
-    from ..smartdevice import SmartDevice
-
 
 class AlarmModule(SmartModule):
     """Implementation of alarm module."""
@@ -23,8 +18,12 @@ def query(self) -> dict:
             "get_support_alarm_type_list": None,  # This should be needed only once
         }
 
-    def __init__(self, device: SmartDevice, module: str):
-        super().__init__(device, module)
+    def _initialize_features(self):
+        """Initialize features.
+
+        This is implemented as some features depend on device responses.
+        """
+        device = self._device
         self._add_feature(
             Feature(
                 device,
@@ -46,12 +45,26 @@ def __init__(self, device: SmartDevice, module: str):
         )
         self._add_feature(
             Feature(
-                device, "Alarm sound", container=self, attribute_getter="alarm_sound"
+                device,
+                "Alarm sound",
+                container=self,
+                attribute_getter="alarm_sound",
+                attribute_setter="set_alarm_sound",
+                category=Feature.Category.Config,
+                type=Feature.Type.Choice,
+                choices_getter="alarm_sounds",
             )
         )
         self._add_feature(
             Feature(
-                device, "Alarm volume", container=self, attribute_getter="alarm_volume"
+                device,
+                "Alarm volume",
+                container=self,
+                attribute_getter="alarm_volume",
+                attribute_setter="set_alarm_volume",
+                category=Feature.Category.Config,
+                type=Feature.Type.Choice,
+                choices=["low", "high"],
             )
         )
         self._add_feature(
@@ -78,6 +91,15 @@ def alarm_sound(self):
         """Return current alarm sound."""
         return self.data["get_alarm_configure"]["type"]
 
+    async def set_alarm_sound(self, sound: str):
+        """Set alarm sound.
+
+        See *alarm_sounds* for list of available sounds.
+        """
+        payload = self.data["get_alarm_configure"].copy()
+        payload["type"] = sound
+        return await self.call("set_alarm_configure", payload)
+
     @property
     def alarm_sounds(self) -> list[str]:
         """Return list of available alarm sounds."""
@@ -88,6 +110,12 @@ def alarm_volume(self):
         """Return alarm volume."""
         return self.data["get_alarm_configure"]["volume"]
 
+    async def set_alarm_volume(self, volume: str):
+        """Set alarm volume."""
+        payload = self.data["get_alarm_configure"].copy()
+        payload["volume"] = volume
+        return await self.call("set_alarm_configure", payload)
+
     @property
     def active(self) -> bool:
         """Return true if alarm is active."""

diff --git a/kasa/smart/smartdevice.py b/kasa/smart/smartdevice.py
@@ -305,6 +305,7 @@ async def _initialize_features(self):
             )
 
         for module in self._modules.values():
+            module._initialize_features()
             for feat in module._module_features.values():
                 self._add_feature(feat)
 

diff --git a/kasa/tests/discovery_fixtures.py b/kasa/tests/discovery_fixtures.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 from dataclasses import dataclass
 from json import dumps as json_dumps
 
@@ -8,7 +9,7 @@
 from kasa.xortransport import XorEncryption
 
 from .fakeprotocol_iot import FakeIotProtocol
-from .fakeprotocol_smart import FakeSmartProtocol
+from .fakeprotocol_smart import FakeSmartProtocol, FakeSmartTransport
 from .fixtureinfo import FixtureInfo, filter_fixtures, idgenerator
 
 
@@ -65,6 +66,7 @@ def parametrize_discovery(desc, *, data_root_filter, protocol_filter=None):
     ids=idgenerator,
 )
 def discovery_mock(request, mocker):
+    """Mock discovery and patch protocol queries to use Fake protocols."""
     fixture_info: FixtureInfo = request.param
     fixture_data = fixture_info.data
 
@@ -157,12 +159,23 @@ async def _query(request, retry_count: int = 3):
 def discovery_data(request, mocker):
     """Return raw discovery file contents as JSON. Used for discovery tests."""
     fixture_info = request.param
-    mocker.patch("kasa.IotProtocol.query", return_value=fixture_info.data)
-    mocker.patch("kasa.SmartProtocol.query", return_value=fixture_info.data)
-    if "discovery_result" in fixture_info.data:
-        return {"result": fixture_info.data["discovery_result"]}
+    fixture_data = copy.deepcopy(fixture_info.data)
+    # Add missing queries to fixture data
+    if "component_nego" in fixture_data:
+        components = {
+            comp["id"]: int(comp["ver_code"])
+            for comp in fixture_data["component_nego"]["component_list"]
+        }
+        for k, v in FakeSmartTransport.FIXTURE_MISSING_MAP.items():
+            # Value is a tuple of component,reponse
+            if k not in fixture_data and v[0] in components:
+                fixture_data[k] = v[1]
+    mocker.patch("kasa.IotProtocol.query", return_value=fixture_data)
+    mocker.patch("kasa.SmartProtocol.query", return_value=fixture_data)
+    if "discovery_result" in fixture_data:
+        return {"result": fixture_data["discovery_result"]}
     else:
-        return {"system": {"get_sysinfo": fixture_info.data["system"]["get_sysinfo"]}}
+        return {"system": {"get_sysinfo": fixture_data["system"]["get_sysinfo"]}}
 
 
 @pytest.fixture(params=UNSUPPORTED_DEVICES.values(), ids=UNSUPPORTED_DEVICES.keys())

diff --git a/kasa/tests/test_cli.py b/kasa/tests/test_cli.py
@@ -354,9 +354,6 @@ async def _state(dev: Device):
 
     mocker.patch("kasa.cli.state", new=_state)
 
-    mocker.patch("kasa.IotProtocol.query", return_value=discovery_mock.query_data)
-    mocker.patch("kasa.SmartProtocol.query", return_value=discovery_mock.query_data)
-
     dr = DiscoveryResult(**discovery_mock.discovery_data["result"])
     res = await runner.invoke(
         cli,

diff --git a/kasa/tests/test_feature.py b/kasa/tests/test_feature.py
@@ -1,4 +1,5 @@
 import pytest
+from pytest_mock import MockFixture
 
 from kasa import Feature
 
@@ -110,6 +111,23 @@ async def test_feature_action(mocker):
     mock_call_action.assert_called()
 
 
+async def test_feature_choice_list(dummy_feature, caplog, mocker: MockFixture):
+    """Test the choice feature type."""
+    dummy_feature.type = Feature.Type.Choice
+    dummy_feature.choices = ["first", "second"]
+
+    mock_setter = mocker.patch.object(dummy_feature.device, "dummysetter", create=True)
+    await dummy_feature.set_value("first")
+    mock_setter.assert_called_with("first")
+    mock_setter.reset_mock()
+
+    with pytest.raises(ValueError):
+        await dummy_feature.set_value("invalid")
+        assert "Unexpected value" in caplog.text
+
+    mock_setter.assert_not_called()
+
+
 @pytest.mark.parametrize("precision_hint", [1, 2, 3])
 async def test_precision_hint(dummy_feature, precision_hint):
     """Test that precision hint works as expected."""