Allow erroring modules to recover (#1080)

Re-query failed modules after some delay instead of immediately disabling them.
Changes to features so they can still be created when modules are erroring.

kasa/smart/modules/alarm.py

@@ -69,7 +69,7 @@ class Alarm(SmartModule):
                 attribute_setter="set_alarm_volume",
                 category=Feature.Category.Config,
                 type=Feature.Type.Choice,
-                choices=["low", "normal", "high"],
+                choices_getter=lambda: ["low", "normal", "high"],
             )
         )
         self._add_feature(

kasa/smart/modules/autooff.py

@@ -39,7 +39,7 @@ class AutoOff(SmartModule):
                 attribute_getter="delay",
                 attribute_setter="set_delay",
                 type=Feature.Type.Number,
-                unit="min",  # ha-friendly unit, see UnitOfTime.MINUTES
+                unit_getter=lambda: "min",  # ha-friendly unit, see UnitOfTime.MINUTES
             )
         )
         self._add_feature(

kasa/smart/modules/batterysensor.py

@@ -37,7 +37,7 @@ class BatterySensor(SmartModule):
                     container=self,
                     attribute_getter="battery",
                     icon="mdi:battery",
-                    unit="%",
+                    unit_getter=lambda: "%",
                     category=Feature.Category.Info,
                     type=Feature.Type.Sensor,
                 )

kasa/smart/modules/brightness.py

@@ -27,8 +27,7 @@ class Brightness(SmartModule):
                 container=self,
                 attribute_getter="brightness",
                 attribute_setter="set_brightness",
-                minimum_value=BRIGHTNESS_MIN,
-                maximum_value=BRIGHTNESS_MAX,
+                range_getter=lambda: (BRIGHTNESS_MIN, BRIGHTNESS_MAX),
                 type=Feature.Type.Number,
                 category=Feature.Category.Primary,
             )

kasa/smart/modules/cloud.py

@@ -18,13 +18,6 @@ class Cloud(SmartModule):
     REQUIRED_COMPONENT = "cloud_connect"
     MINIMUM_UPDATE_INTERVAL_SECS = 60
 
-    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)
 

kasa/smart/modules/energy.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 from ...emeterstatus import EmeterStatus
 from ...exceptions import KasaException
 from ...interfaces.energy import Energy as EnergyInterface
-from ..smartmodule import SmartModule
+from ..smartmodule import SmartModule, raise_if_update_error
 
 
 class Energy(SmartModule, EnergyInterface):
@@ -23,6 +23,7 @@ class Energy(SmartModule, EnergyInterface):
         return req
 
     @property
+    @raise_if_update_error
     def current_consumption(self) -> float | None:
         """Current power in watts."""
         if (power := self.energy.get("current_power")) is not None:
@@ -30,6 +31,7 @@ class Energy(SmartModule, EnergyInterface):
         return None
 
     @property
+    @raise_if_update_error
     def energy(self):
         """Return get_energy_usage results."""
         if en := self.data.get("get_energy_usage"):
@@ -45,6 +47,7 @@ class Energy(SmartModule, EnergyInterface):
         )
 
     @property
+    @raise_if_update_error
     def status(self):
         """Get the emeter status."""
         return self._get_status_from_energy(self.energy)
@@ -55,26 +58,31 @@ class Energy(SmartModule, EnergyInterface):
         return self._get_status_from_energy(res["get_energy_usage"])
 
     @property
+    @raise_if_update_error
     def consumption_this_month(self) -> float | None:
         """Get the emeter value for this month in kWh."""
         return self.energy.get("month_energy") / 1_000
 
     @property
+    @raise_if_update_error
     def consumption_today(self) -> float | None:
         """Get the emeter value for today in kWh."""
         return self.energy.get("today_energy") / 1_000
 
     @property
+    @raise_if_update_error
     def consumption_total(self) -> float | None:
         """Return total consumption since last reboot in kWh."""
         return None
 
     @property
+    @raise_if_update_error
     def current(self) -> float | None:
         """Return the current in A."""
         return None
 
     @property
+    @raise_if_update_error
     def voltage(self) -> float | None:
         """Get the current voltage in V."""
         return None

kasa/smart/modules/fan.py

@@ -30,8 +30,7 @@ class Fan(SmartModule, FanInterface):
                 attribute_setter="set_fan_speed_level",
                 icon="mdi:fan",
                 type=Feature.Type.Number,
-                minimum_value=0,
-                maximum_value=4,
+                range_getter=lambda: (0, 4),
                 category=Feature.Category.Primary,
             )
         )

kasa/smart/modules/humiditysensor.py

@@ -27,7 +27,7 @@ class HumiditySensor(SmartModule):
                 container=self,
                 attribute_getter="humidity",
                 icon="mdi:water-percent",
-                unit="%",
+                unit_getter=lambda: "%",
                 category=Feature.Category.Primary,
                 type=Feature.Type.Sensor,
             )

kasa/smart/modules/lighttransition.py

@@ -73,7 +73,7 @@ class LightTransition(SmartModule):
                     attribute_setter="set_turn_on_transition",
                     icon=icon,
                     type=Feature.Type.Number,
-                    maximum_value=self._turn_on_transition_max,
+                    range_getter=lambda: (0, self._turn_on_transition_max),
                 )
             )
             self._add_feature(
@@ -86,7 +86,7 @@ class LightTransition(SmartModule):
                     attribute_setter="set_turn_off_transition",
                     icon=icon,
                     type=Feature.Type.Number,
-                    maximum_value=self._turn_off_transition_max,
+                    range_getter=lambda: (0, self._turn_off_transition_max),
                 )
             )
 

kasa/smart/modules/reportmode.py

@@ -26,7 +26,7 @@ class ReportMode(SmartModule):
                 name="Report interval",
                 container=self,
                 attribute_getter="report_interval",
-                unit="s",
+                unit_getter=lambda: "s",
                 category=Feature.Category.Debug,
                 type=Feature.Type.Sensor,
             )

kasa/smart/modules/temperaturecontrol.py

@@ -51,8 +51,7 @@ class TemperatureControl(SmartModule):
                 container=self,
                 attribute_getter="temperature_offset",
                 attribute_setter="set_temperature_offset",
-                minimum_value=-10,
-                maximum_value=10,
+                range_getter=lambda: (-10, 10),
                 type=Feature.Type.Number,
                 category=Feature.Category.Config,
             )

kasa/smart/modules/temperaturesensor.py

@@ -54,7 +54,7 @@ class TemperatureSensor(SmartModule):
                 attribute_getter="temperature_unit",
                 attribute_setter="set_temperature_unit",
                 type=Feature.Type.Choice,
-                choices=["celsius", "fahrenheit"],
+                choices_getter=lambda: ["celsius", "fahrenheit"],
             )
         )
 

kasa/smart/smartchilddevice.py

@@ -10,6 +10,7 @@ from ..device_type import DeviceType
 from ..deviceconfig import DeviceConfig
 from ..smartprotocol import SmartProtocol, _ChildProtocolWrapper
 from .smartdevice import SmartDevice
+from .smartmodule import SmartModule
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -49,13 +50,21 @@ class SmartChildDevice(SmartDevice):
         Internal implementation to allow patching of public update in the cli
         or test framework.
         """
+        now = time.monotonic()
+        module_queries: list[SmartModule] = []
         req: dict[str, Any] = {}
         for module in self.modules.values():
-            if mod_query := module.query():
+            if module.disabled is False and (mod_query := module.query()):
+                module_queries.append(module)
                 req.update(mod_query)
         if req:
             self._last_update = await self.protocol.query(req)
-            self._last_update_time = time.time()
+
+        for module in self.modules.values():
+            self._handle_module_post_update(
+                module, now, had_query=module in module_queries
+            )
+        self._last_update_time = now
 
     @classmethod
     async def create(cls, parent: SmartDevice, child_info, child_components):

kasa/smart/smartdevice.py

@@ -165,28 +165,25 @@ class SmartDevice(Device):
         if first_update:
             await self._negotiate()
             await self._initialize_modules()
+            # Run post update for the cloud module
+            if cloud_mod := self.modules.get(Module.Cloud):
+                self._handle_module_post_update(cloud_mod, now, had_query=True)
 
         resp = await self._modular_update(first_update, now)
 
-        # Call child update which will only update module calls, info is updated
-        # from get_child_device_list. update_children only affects hub devices, other
-        # devices will always update children to prevent errors on module access.
-        if update_children or self.device_type != DeviceType.Hub:
-            for child in self._children.values():
-                await child._update()
         if child_info := self._try_get_response(
             self._last_update, "get_child_device_list", {}
         ):
             for info in child_info["child_device_list"]:
                 self._children[info["device_id"]]._update_internal_state(info)
-
-        for child in self._children.values():
-            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)
+        # Call child update which will only update module calls, info is updated
+        # from get_child_device_list. update_children only affects hub devices, other
+        # devices will always update children to prevent errors on module access.
+        # This needs to go after updating the internal state of the children so that
+        # child modules have access to their sysinfo.
+        if update_children or self.device_type != DeviceType.Hub:
+            for child in self._children.values():
+                await child._update()
 
         # We can first initialize the features after the first update.
         # We make here an assumption that every device has at least a single feature.
@@ -197,18 +194,26 @@ class SmartDevice(Device):
             updated = self._last_update if first_update else resp
             _LOGGER.debug("Update completed %s: %s", self.host, list(updated.keys()))
 
-    def _handle_module_post_update_hook(self, module: SmartModule) -> bool:
+    def _handle_module_post_update(
+        self, module: SmartModule, update_time: float, had_query: bool
+    ):
+        if module.disabled:
+            return  # pragma: no cover
+        if had_query:
+            module._last_update_time = update_time
         try:
             module._post_update_hook()
-            return True
+            module._set_error(None)
         except Exception as ex:
-            _LOGGER.warning(
-                "Error processing %s for device %s, module will be unavailable: %s",
-                module.name,
-                self.host,
-                ex,
-            )
-            return False
+            # Only set the error if a query happened.
+            if had_query:
+                module._set_error(ex)
+                _LOGGER.warning(
+                    "Error processing %s for device %s, module will be unavailable: %s",
+                    module.name,
+                    self.host,
+                    ex,
+                )
 
     async def _modular_update(
         self, first_update: bool, update_time: float
@@ -221,17 +226,16 @@ class SmartDevice(Device):
         mq = {
             module: query
             for module in self._modules.values()
-            if (query := module.query())
+            if module.disabled is False and (query := module.query())
         }
         for module, query in mq.items():
             if first_update and module.__class__ in FIRST_UPDATE_MODULES:
                 module._last_update_time = update_time
                 continue
             if (
-                not module.MINIMUM_UPDATE_INTERVAL_SECS
+                not module.update_interval
                 or not module._last_update_time
-                or (update_time - module._last_update_time)
-                >= module.MINIMUM_UPDATE_INTERVAL_SECS
+                or (update_time - module._last_update_time) >= module.update_interval
             ):
                 module_queries.append(module)
                 req.update(query)
@@ -254,16 +258,10 @@ class SmartDevice(Device):
         self._info = self._try_get_response(info_resp, "get_device_info")
 
         # Call handle update for modules that want to update internal data
-        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)
-
-        # Set the last update time for modules that had queries made.
-        for module in module_queries:
-            module._last_update_time = update_time
+        for module in self._modules.values():
+            self._handle_module_post_update(
+                module, update_time, had_query=module in module_queries
+            )
 
         return resp
 
@@ -392,7 +390,7 @@ class SmartDevice(Device):
                     name="RSSI",
                     attribute_getter=lambda x: x._info["rssi"],
                     icon="mdi:signal",
-                    unit="dBm",
+                    unit_getter=lambda: "dBm",
                     category=Feature.Category.Debug,
                     type=Feature.Type.Sensor,
                 )

kasa/smart/smartmodule.py

@@ -18,6 +18,7 @@ _LOGGER = logging.getLogger(__name__)
 
 _T = TypeVar("_T", bound="SmartModule")
 _P = ParamSpec("_P")
+_R = TypeVar("_R")
 
 
 def allow_update_after(
@@ -38,6 +39,17 @@ def allow_update_after(
     return _async_wrap
 
 
+def raise_if_update_error(func: Callable[[_T], _R]) -> Callable[[_T], _R]:
+    """Define a wrapper to raise an error if the last module update was an error."""
+
+    def _wrap(self: _T) -> _R:
+        if err := self._last_update_error:
+            raise err
+        return func(self)
+
+    return _wrap
+
+
 class SmartModule(Module):
     """Base class for SMART modules."""
 
@@ -52,17 +64,58 @@ class SmartModule(Module):
     REGISTERED_MODULES: dict[str, type[SmartModule]] = {}
 
     MINIMUM_UPDATE_INTERVAL_SECS = 0
+    UPDATE_INTERVAL_AFTER_ERROR_SECS = 30
+
+    DISABLE_AFTER_ERROR_COUNT = 10
 
     def __init__(self, device: SmartDevice, module: str):
         self._device: SmartDevice
         super().__init__(device, module)
         self._last_update_time: float | None = None
+        self._last_update_error: KasaException | None = None
+        self._error_count = 0
 
     def __init_subclass__(cls, **kwargs):
         name = getattr(cls, "NAME", cls.__name__)
         _LOGGER.debug("Registering %s" % cls)
         cls.REGISTERED_MODULES[name] = cls
 
+    def _set_error(self, err: Exception | None):
+        if err is None:
+            self._error_count = 0
+            self._last_update_error = None
+        else:
+            self._last_update_error = KasaException("Module update error", err)
+            self._error_count += 1
+            if self._error_count == self.DISABLE_AFTER_ERROR_COUNT:
+                _LOGGER.error(
+                    "Error processing %s for device %s, module will be disabled: %s",
+                    self.name,
+                    self._device.host,
+                    err,
+                )
+            if self._error_count > self.DISABLE_AFTER_ERROR_COUNT:
+                _LOGGER.error(  # pragma: no cover
+                    "Unexpected error processing %s for device %s, "
+                    "module should be disabled: %s",
+                    self.name,
+                    self._device.host,
+                    err,
+                )
+
+    @property
+    def update_interval(self) -> int:
+        """Time to wait between updates."""
+        if self._last_update_error is None:
+            return self.MINIMUM_UPDATE_INTERVAL_SECS
+
+        return self.UPDATE_INTERVAL_AFTER_ERROR_SECS * self._error_count
+
+    @property
+    def disabled(self) -> bool:
+        """Return true if the module is disabled due to errors."""
+        return self._error_count >= self.DISABLE_AFTER_ERROR_COUNT
+
     @property
     def name(self) -> str:
         """Name of the module."""