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

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.
[#1033](https://github.com/python-kasa/python-kasa/issues/1033)

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,

kasa/smart/modules/batterysensor.py

@@ -43,6 +43,10 @@ class BatterySensor(SmartModule):
                 )
             )
 
+    def query(self) -> dict:
+        """Query to execute during the update cycle."""
+        return {}
+
     @property
     def battery(self):
         """Return battery level."""

kasa/smart/modules/cloud.py

@@ -4,7 +4,6 @@ from __future__ import annotations
 
 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 @@ class Cloud(SmartModule):
     @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

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 = {

kasa/smart/modules/firmware.py

@@ -13,7 +13,6 @@ from typing import TYPE_CHECKING, Any, Callable, Optional
 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 @@ class Firmware(SmartModule):
             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 @@ class Firmware(SmartModule):
     @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

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."""

kasa/smart/modules/humiditysensor.py

@@ -45,6 +45,10 @@ class HumiditySensor(SmartModule):
             )
         )
 
+    def query(self) -> dict:
+        """Query to execute during the update cycle."""
+        return {}
+
     @property
     def humidity(self):
         """Return current humidity in percentage."""

kasa/smart/modules/lightpreset.py

@@ -140,7 +140,7 @@ class LightPreset(SmartModule, LightPresetInterface):
         """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.

kasa/smart/modules/lighttransition.py

@@ -230,7 +230,7 @@ class LightTransition(SmartModule):
         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."""

kasa/smart/modules/reportmode.py

@@ -32,6 +32,10 @@ class ReportMode(SmartModule):
             )
         )
 
+    def query(self) -> dict:
+        """Query to execute during the update cycle."""
+        return {}
+
     @property
     def report_interval(self):
         """Reporting interval of a sensor device."""

kasa/smart/modules/temperaturesensor.py

@@ -58,6 +58,10 @@ class TemperatureSensor(SmartModule):
             )
         )
 
+    def query(self) -> dict:
+        """Query to execute during the update cycle."""
+        return {}
+
     @property
     def temperature(self):
         """Return current humidity in percentage."""