From adce92a761e899b7751b8667dfeef9f80561f7c4 Mon Sep 17 00:00:00 2001
From: Steven B <51370195+sdb9696@users.noreply.github.com>
Date: Wed, 6 Mar 2024 16:45:08 +0000
Subject: [PATCH] Add iot brightness feature (#808)

---
 kasa/iot/iotbulb.py                   | 17 +++++++++++++++++
 kasa/iot/iotdimmer.py                 | 17 +++++++++++++++++
 kasa/iot/iotplug.py                   |  3 +++
 kasa/tests/test_feature_brightness.py | 25 ++++++++++++++++++++++++-
 4 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/kasa/iot/iotbulb.py b/kasa/iot/iotbulb.py
index 6b8d37b0..d80a24ea 100644
--- a/kasa/iot/iotbulb.py
+++ b/kasa/iot/iotbulb.py
@@ -12,6 +12,7 @@ except ImportError:
 from ..bulb import HSV, Bulb, BulbPreset, ColorTempRange
 from ..device_type import DeviceType
 from ..deviceconfig import DeviceConfig
+from ..feature import Feature, FeatureType
 from ..protocol import BaseProtocol
 from .iotdevice import IotDevice, KasaException, requires_update
 from .modules import Antitheft, Cloud, Countdown, Emeter, Schedule, Time, Usage
@@ -204,6 +205,22 @@ class IotBulb(IotDevice, Bulb):
         self.add_module("countdown", Countdown(self, "countdown"))
         self.add_module("cloud", Cloud(self, "smartlife.iot.common.cloud"))
 
+    async def _initialize_features(self):
+        await super()._initialize_features()
+
+        if bool(self.sys_info["is_dimmable"]):  # pragma: no branch
+            self._add_feature(
+                Feature(
+                    device=self,
+                    name="Brightness",
+                    attribute_getter="brightness",
+                    attribute_setter="set_brightness",
+                    minimum_value=1,
+                    maximum_value=100,
+                    type=FeatureType.Number,
+                )
+            )
+
     @property  # type: ignore
     @requires_update
     def is_color(self) -> bool:
diff --git a/kasa/iot/iotdimmer.py b/kasa/iot/iotdimmer.py
index 721a2c4b..8882ae81 100644
--- a/kasa/iot/iotdimmer.py
+++ b/kasa/iot/iotdimmer.py
@@ -4,6 +4,7 @@ from typing import Any, Dict, Optional
 
 from ..device_type import DeviceType
 from ..deviceconfig import DeviceConfig
+from ..feature import Feature, FeatureType
 from ..protocol import BaseProtocol
 from .iotdevice import KasaException, requires_update
 from .iotplug import IotPlug
@@ -80,6 +81,22 @@ class IotDimmer(IotPlug):
         self.add_module("motion", Motion(self, "smartlife.iot.PIR"))
         self.add_module("ambient", AmbientLight(self, "smartlife.iot.LAS"))
 
+    async def _initialize_features(self):
+        await super()._initialize_features()
+
+        if "brightness" in self.sys_info:  # pragma: no branch
+            self._add_feature(
+                Feature(
+                    device=self,
+                    name="Brightness",
+                    attribute_getter="brightness",
+                    attribute_setter="set_brightness",
+                    minimum_value=1,
+                    maximum_value=100,
+                    type=FeatureType.Number,
+                )
+            )
+
     @property  # type: ignore
     @requires_update
     def brightness(self) -> int:
diff --git a/kasa/iot/iotplug.py b/kasa/iot/iotplug.py
index 3f776b98..2d509e05 100644
--- a/kasa/iot/iotplug.py
+++ b/kasa/iot/iotplug.py
@@ -57,6 +57,9 @@ class IotPlug(IotDevice):
         self.add_module("time", Time(self, "time"))
         self.add_module("cloud", Cloud(self, "cnCloud"))
 
+    async def _initialize_features(self):
+        await super()._initialize_features()
+
         self._add_feature(
             Feature(
                 device=self,
diff --git a/kasa/tests/test_feature_brightness.py b/kasa/tests/test_feature_brightness.py
index 9d9d3165..72bc3637 100644
--- a/kasa/tests/test_feature_brightness.py
+++ b/kasa/tests/test_feature_brightness.py
@@ -1,7 +1,8 @@
 import pytest
 
+from kasa.iot import IotDevice
 from kasa.smart import SmartDevice
-from kasa.tests.conftest import parametrize
+from kasa.tests.conftest import dimmable, parametrize
 
 brightness = parametrize("brightness smart", component_filter="brightness")
 
@@ -26,3 +27,25 @@ async def test_brightness_component(dev: SmartDevice):
 
     with pytest.raises(ValueError):
         await feature.set_value(feature.maximum_value + 10)
+
+
+@dimmable
+async def test_brightness_dimmable(dev: SmartDevice):
+    """Test brightness feature."""
+    assert isinstance(dev, IotDevice)
+    assert "brightness" in dev.sys_info or bool(dev.sys_info["is_dimmable"])
+
+    # Test getting the value
+    feature = dev.features["brightness"]
+    assert isinstance(feature.value, int)
+    assert feature.value > 0 and feature.value <= 100
+
+    # Test setting the value
+    await feature.set_value(10)
+    assert feature.value == 10
+
+    with pytest.raises(ValueError):
+        await feature.set_value(feature.minimum_value - 10)
+
+    with pytest.raises(ValueError):
+        await feature.set_value(feature.maximum_value + 10)