From 0d5a3c84391f903236b54fe48698076d39e74703 Mon Sep 17 00:00:00 2001
From: Teemu R <tpr@iki.fi>
Date: Tue, 5 Mar 2024 15:41:40 +0100
Subject: [PATCH] Add brightness module (#806)

Add module for controlling the brightness.
---
 kasa/smart/modules/__init__.py        |  2 ++
 kasa/smart/modules/brightness.py      | 43 +++++++++++++++++++++++++++
 kasa/smart/smartmodule.py             | 11 +++++--
 kasa/tests/device_fixtures.py         |  2 --
 kasa/tests/test_feature_brightness.py | 26 ++++++++++++----
 5 files changed, 74 insertions(+), 10 deletions(-)
 create mode 100644 kasa/smart/modules/brightness.py

diff --git a/kasa/smart/modules/__init__.py b/kasa/smart/modules/__init__.py
index 3e95dfe7..dc4e0cf5 100644
--- a/kasa/smart/modules/__init__.py
+++ b/kasa/smart/modules/__init__.py
@@ -2,6 +2,7 @@
 from .alarmmodule import AlarmModule
 from .autooffmodule import AutoOffModule
 from .battery import BatterySensor
+from .brightness import Brightness
 from .childdevicemodule import ChildDeviceModule
 from .cloudmodule import CloudModule
 from .devicemodule import DeviceModule
@@ -26,6 +27,7 @@ __all__ = [
     "ReportModule",
     "AutoOffModule",
     "LedModule",
+    "Brightness",
     "Firmware",
     "CloudModule",
     "LightTransitionModule",
diff --git a/kasa/smart/modules/brightness.py b/kasa/smart/modules/brightness.py
new file mode 100644
index 00000000..03e9e238
--- /dev/null
+++ b/kasa/smart/modules/brightness.py
@@ -0,0 +1,43 @@
+"""Implementation of brightness module."""
+from typing import TYPE_CHECKING, Dict
+
+from ...feature import Feature, FeatureType
+from ..smartmodule import SmartModule
+
+if TYPE_CHECKING:
+    from ..smartdevice import SmartDevice
+
+
+class Brightness(SmartModule):
+    """Implementation of brightness module."""
+
+    REQUIRED_COMPONENT = "brightness"
+
+    def __init__(self, device: "SmartDevice", module: str):
+        super().__init__(device, module)
+        self._add_feature(
+            Feature(
+                device,
+                "Brightness",
+                container=self,
+                attribute_getter="brightness",
+                attribute_setter="set_brightness",
+                minimum_value=1,
+                maximum_value=100,
+                type=FeatureType.Number,
+            )
+        )
+
+    def query(self) -> Dict:
+        """Query to execute during the update cycle."""
+        # Brightness is contained in the main device info response.
+        return {}
+
+    @property
+    def brightness(self):
+        """Return current brightness."""
+        return self.data["brightness"]
+
+    async def set_brightness(self, brightness: int):
+        """Set the brightness."""
+        return await self.call("set_device_info", {"brightness": brightness})
diff --git a/kasa/smart/smartmodule.py b/kasa/smart/smartmodule.py
index e34f2260..01a27360 100644
--- a/kasa/smart/smartmodule.py
+++ b/kasa/smart/smartmodule.py
@@ -53,14 +53,19 @@ class SmartModule(Module):
     def data(self):
         """Return response data for the module.
 
-        If module performs only a single query, the resulting response is unwrapped.
+        If the module performs only a single query, the resulting response is unwrapped.
+        If the module does not define a query, this property returns a reference
+        to the main "get_device_info" response.
         """
+        dev = self._device
         q = self.query()
+
+        if not q:
+            return dev.internal_state["get_device_info"]
+
         q_keys = list(q.keys())
         query_key = q_keys[0]
 
-        dev = self._device
-
         # TODO: hacky way to check if update has been called.
         #  The way this falls back to parent may not always be wanted.
         #  Especially, devices can have their own firmware updates.
diff --git a/kasa/tests/device_fixtures.py b/kasa/tests/device_fixtures.py
index 73d171d2..8a1b643b 100644
--- a/kasa/tests/device_fixtures.py
+++ b/kasa/tests/device_fixtures.py
@@ -249,8 +249,6 @@ device_iot = parametrize(
     "devices iot", model_filter=ALL_DEVICES_IOT, protocol_filter={"IOT"}
 )
 
-brightness = parametrize("brightness smart", component_filter="brightness")
-
 
 def check_categories():
     """Check that every fixture file is categorized."""
diff --git a/kasa/tests/test_feature_brightness.py b/kasa/tests/test_feature_brightness.py
index d99b55d1..9d9d3165 100644
--- a/kasa/tests/test_feature_brightness.py
+++ b/kasa/tests/test_feature_brightness.py
@@ -1,12 +1,28 @@
-from kasa.smart import SmartDevice
+import pytest
 
-from .conftest import (
-    brightness,
-)
+from kasa.smart import SmartDevice
+from kasa.tests.conftest import parametrize
+
+brightness = parametrize("brightness smart", component_filter="brightness")
 
 
 @brightness
 async def test_brightness_component(dev: SmartDevice):
-    """Placeholder to test framwework component filter."""
+    """Test brightness feature."""
     assert isinstance(dev, SmartDevice)
     assert "brightness" in dev._components
+
+    # 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)