Handle module errors more robustly and add query params to light preset and transition (#1043)

sdb9696 · web-flow · commit fe116eaefbe2 · 2024-07-04T08:29:53.000+01:00
Ensures that all modules try to access their data in `_post_update_hook` in a safe manner and disable themselves if there's an error. Also adds parameters to get_preset_rules and get_on_off_gradually_info to fix issues with recent firmware updates. Cherry pick of [#1036](#1036) to patch
diff --git a/devtools/helpers/smartrequests.py b/devtools/helpers/smartrequests.py
@@ -284,6 +284,15 @@ def get_preset_rules(params: GetRulesParams | None = None) -> SmartRequest:
         """Get preset rules."""
         return SmartRequest("get_preset_rules", params or SmartRequest.GetRulesParams())
 
+    @staticmethod
+    def get_on_off_gradually_info(
+        params: SmartRequestParams | None = None,
+    ) -> SmartRequest:
+        """Get preset rules."""
+        return SmartRequest(
+            "get_on_off_gradually_info", params or SmartRequest.SmartRequestParams()
+        )
+
     @staticmethod
     def get_auto_light_info() -> SmartRequest:
         """Get auto light info."""
@@ -382,7 +391,7 @@ def get_component_requests(component_id, ver_code):
     "auto_light": [SmartRequest.get_auto_light_info()],
     "light_effect": [SmartRequest.get_dynamic_light_effect_rules()],
     "bulb_quick_control": [],
-    "on_off_gradually": [SmartRequest.get_raw_request("get_on_off_gradually_info")],
+    "on_off_gradually": [SmartRequest.get_on_off_gradually_info()],
     "light_strip": [],
     "light_strip_lighting_effect": [
         SmartRequest.get_raw_request("get_lighting_effect")
diff --git a/kasa/smart/modules/autooff.py b/kasa/smart/modules/autooff.py
@@ -19,12 +19,6 @@ class AutoOff(SmartModule):
 
     def _initialize_features(self):
         """Initialize features after the initial update."""
-        if not isinstance(self.data, dict):
-            _LOGGER.warning(
-                "No data available for module, skipping %s: %s", self, self.data
-            )
-            return
-
         self._add_feature(
             Feature(
                 self._device,
diff --git a/kasa/smart/modules/batterysensor.py b/kasa/smart/modules/batterysensor.py
@@ -43,6 +43,10 @@ def _initialize_features(self):
                 )
             )
 
+    def query(self) -> dict:
+        """Query to execute during the update cycle."""
+        return {}
+
     @property
     def battery(self):
         """Return battery level."""
diff --git a/kasa/smart/modules/cloud.py b/kasa/smart/modules/cloud.py
@@ -4,7 +4,6 @@
 
 from typing import TYPE_CHECKING
 
-from ...exceptions import SmartErrorCode
 from ...feature import Feature
 from ..smartmodule import SmartModule
 
@@ -18,6 +17,13 @@ class Cloud(SmartModule):
     QUERY_GETTER_NAME = "get_connect_cloud_state"
     REQUIRED_COMPONENT = "cloud_connect"
 
+    def _post_update_hook(self):
+        """Perform actions after a device update.
+
+        Overrides the default behaviour to disable a module if the query returns
+        an error because the logic here is to treat that as not connected.
+        """
+
     def __init__(self, device: SmartDevice, module: str):
         super().__init__(device, module)
 
@@ -37,6 +43,6 @@ def __init__(self, device: SmartDevice, module: str):
     @property
     def is_connected(self):
         """Return True if device is connected to the cloud."""
-        if isinstance(self.data, SmartErrorCode):
+        if self._has_data_error():
             return False
         return self.data["status"] == 0
diff --git a/kasa/smart/modules/devicemodule.py b/kasa/smart/modules/devicemodule.py
@@ -10,6 +10,13 @@ class DeviceModule(SmartModule):
 
     REQUIRED_COMPONENT = "device"
 
+    def _post_update_hook(self):
+        """Perform actions after a device update.
+
+        Overrides the default behaviour to disable a module if the query returns
+        an error because this module is critical.
+        """
+
     def query(self) -> dict:
         """Query to execute during the update cycle."""
         query = {
diff --git a/kasa/smart/modules/firmware.py b/kasa/smart/modules/firmware.py
@@ -13,7 +13,6 @@
 from async_timeout import timeout as asyncio_timeout
 from pydantic.v1 import BaseModel, Field, validator
 
-from ...exceptions import SmartErrorCode
 from ...feature import Feature
 from ..smartmodule import SmartModule
 
@@ -123,6 +122,13 @@ def query(self) -> dict:
             req["get_auto_update_info"] = None
         return req
 
+    def _post_update_hook(self):
+        """Perform actions after a device update.
+
+        Overrides the default behaviour to disable a module if the query returns
+        an error because some of the module still functions.
+        """
+
     @property
     def current_firmware(self) -> str:
         """Return the current firmware version."""
@@ -136,11 +142,11 @@ def latest_firmware(self) -> str:
     @property
     def firmware_update_info(self):
         """Return latest firmware information."""
-        fw = self.data.get("get_latest_fw") or self.data
-        if not self._device.is_cloud_connected or isinstance(fw, SmartErrorCode):
+        if not self._device.is_cloud_connected or self._has_data_error():
             # Error in response, probably disconnected from the cloud.
             return UpdateInfo(type=0, need_to_upgrade=False)
 
+        fw = self.data.get("get_latest_fw") or self.data
         return UpdateInfo.parse_obj(fw)
 
     @property
diff --git a/kasa/smart/modules/frostprotection.py b/kasa/smart/modules/frostprotection.py
@@ -14,6 +14,10 @@ class FrostProtection(SmartModule):
     REQUIRED_COMPONENT = "frost_protection"
     QUERY_GETTER_NAME = "get_frost_protection"
 
+    def query(self) -> dict:
+        """Query to execute during the update cycle."""
+        return {}
+
     @property
     def enabled(self) -> bool:
         """Return True if frost protection is on."""
diff --git a/kasa/smart/modules/humiditysensor.py b/kasa/smart/modules/humiditysensor.py
@@ -45,6 +45,10 @@ def __init__(self, device: SmartDevice, module: str):
             )
         )
 
+    def query(self) -> dict:
+        """Query to execute during the update cycle."""
+        return {}
+
     @property
     def humidity(self):
         """Return current humidity in percentage."""
diff --git a/kasa/smart/modules/lightpreset.py b/kasa/smart/modules/lightpreset.py
@@ -140,7 +140,7 @@ def query(self) -> dict:
         """Query to execute during the update cycle."""
         if self._state_in_sysinfo:  # Child lights can have states in the child info
             return {}
-        return {self.QUERY_GETTER_NAME: None}
+        return {self.QUERY_GETTER_NAME: {"start_index": 0}}
 
     async def _check_supported(self):
         """Additional check to see if the module is supported by the device.
diff --git a/kasa/smart/modules/lighttransition.py b/kasa/smart/modules/lighttransition.py
@@ -230,7 +230,7 @@ def query(self) -> dict:
         if self._state_in_sysinfo:
             return {}
         else:
-            return {self.QUERY_GETTER_NAME: None}
+            return {self.QUERY_GETTER_NAME: {}}
 
     async def _check_supported(self):
         """Additional check to see if the module is supported by the device."""
diff --git a/kasa/smart/modules/reportmode.py b/kasa/smart/modules/reportmode.py
@@ -32,6 +32,10 @@ def __init__(self, device: SmartDevice, module: str):
             )
         )
 
+    def query(self) -> dict:
+        """Query to execute during the update cycle."""
+        return {}
+
     @property
     def report_interval(self):
         """Reporting interval of a sensor device."""
diff --git a/kasa/smart/modules/temperaturesensor.py b/kasa/smart/modules/temperaturesensor.py
@@ -58,6 +58,10 @@ def __init__(self, device: SmartDevice, module: str):
             )
         )
 
+    def query(self) -> dict:
+        """Query to execute during the update cycle."""
+        return {}
+
     @property
     def temperature(self):
         """Return current humidity in percentage."""
diff --git a/kasa/smart/smartdevice.py b/kasa/smart/smartdevice.py
@@ -177,11 +177,20 @@ async def update(self, update_children: bool = False):
                 self._children[info["device_id"]]._update_internal_state(info)
 
         # Call handle update for modules that want to update internal data
-        for module in self._modules.values():
-            module._post_update_hook()
+        errors = []
+        for module_name, module in self._modules.items():
+            if not self._handle_module_post_update_hook(module):
+                errors.append(module_name)
+        for error in errors:
+            self._modules.pop(error)
+
         for child in self._children.values():
-            for child_module in child._modules.values():
-                child_module._post_update_hook()
+            errors = []
+            for child_module_name, child_module in child._modules.items():
+                if not self._handle_module_post_update_hook(child_module):
+                    errors.append(child_module_name)
+            for error in errors:
+                child._modules.pop(error)
 
         # We can first initialize the features after the first update.
         # We make here an assumption that every device has at least a single feature.
@@ -190,6 +199,19 @@ async def update(self, update_children: bool = False):
 
         _LOGGER.debug("Got an update: %s", self._last_update)
 
+    def _handle_module_post_update_hook(self, module: SmartModule) -> bool:
+        try:
+            module._post_update_hook()
+            return True
+        except Exception as ex:
+            _LOGGER.error(
+                "Error processing %s for device %s, module will be unavailable: %s",
+                module.name,
+                self.host,
+                ex,
+            )
+            return False
+
     async def _initialize_modules(self):
         """Initialize modules based on component negotiation response."""
         from .smartmodule import SmartModule
diff --git a/kasa/smart/smartmodule.py b/kasa/smart/smartmodule.py
@@ -5,7 +5,7 @@
 import logging
 from typing import TYPE_CHECKING
 
-from ..exceptions import KasaException
+from ..exceptions import DeviceError, KasaException, SmartErrorCode
 from ..module import Module
 
 if TYPE_CHECKING:
@@ -41,6 +41,14 @@ def name(self) -> str:
         """Name of the module."""
         return getattr(self, "NAME", self.__class__.__name__)
 
+    def _post_update_hook(self):  # noqa: B027
+        """Perform actions after a device update.
+
+        Any modules overriding this should ensure that self.data is
+        accessed unless the module should remain active despite errors.
+        """
+        assert self.data  # noqa: S101
+
     def query(self) -> dict:
         """Query to execute during the update cycle.
 
@@ -87,6 +95,11 @@ def data(self):
 
         filtered_data = {k: v for k, v in dev._last_update.items() if k in q_keys}
 
+        for data_item in filtered_data:
+            if isinstance(filtered_data[data_item], SmartErrorCode):
+                raise DeviceError(
+                    f"{data_item} for {self.name}", error_code=filtered_data[data_item]
+                )
         if len(filtered_data) == 1:
             return next(iter(filtered_data.values()))
 
@@ -110,3 +123,10 @@ async def _check_supported(self) -> bool:
         color_temp_range but only supports one value.
         """
         return True
+
+    def _has_data_error(self) -> bool:
+        try:
+            assert self.data  # noqa: S101
+            return False
+        except DeviceError:
+            return True
diff --git a/kasa/smartprotocol.py b/kasa/smartprotocol.py
@@ -416,6 +416,10 @@ def _get_method_and_params_for_request(self, request):
         return smart_method, smart_params
 
     async def query(self, request: str | dict, retry_count: int = 3) -> dict:
+        """Wrap request inside control_child envelope."""
+        return await self._query(request, retry_count)
+
+    async def _query(self, request: str | dict, retry_count: int = 3) -> dict:
         """Wrap request inside control_child envelope."""
         method, params = self._get_method_and_params_for_request(request)
         request_data = {
diff --git a/kasa/tests/test_smartdevice.py b/kasa/tests/test_smartdevice.py
diff --git a/kasa/tests/test_smartprotocol.py b/kasa/tests/test_smartprotocol.py

Original file line number	Diff line number	Diff line change
`@@ -43,6 +43,10 @@ def _initialize_features(self):`
`43`	`43`	`)`
`44`	`44`	`)`
`45`	`45`
	`46`	`+ def query(self) -> dict:`
	`47`	`+ """Query to execute during the update cycle."""`
	`48`	`+ return {}`
	`49`	`+`
`46`	`50`	`@property`
`47`	`51`	`def battery(self):`
`48`	`52`	`"""Return battery level."""`
Original file line number	Diff line number	Diff line change
`@@ -45,6 +45,10 @@ def __init__(self, device: SmartDevice, module: str):`
`45`	`45`	`)`
`46`	`46`	`)`
`47`	`47`
	`48`	`+ def query(self) -> dict:`
	`49`	`+ """Query to execute during the update cycle."""`
	`50`	`+ return {}`
	`51`	`+`
`48`	`52`	`@property`
`49`	`53`	`def humidity(self):`
`50`	`54`	`"""Return current humidity in percentage."""`