Fallback to other module data on get_energy_usage errors (#1245)

- The `get_energy_usage` query can fail if the device time is not set because the response includes the device time.
- Make `get_energy_usage` an optional query response so the energy module can fall back to getting the power from `get_emeter_data` or `get_current_power` on error.
- Devices on `energy_monitoring` version 1 still fail as they have no additional queries to fall back to.

kasa/smart/modules/energy.py

@@ -2,10 +2,10 @@
 
 from __future__ import annotations
 
-from typing import NoReturn
+from typing import Any, NoReturn
 
 from ...emeterstatus import EmeterStatus
-from ...exceptions import KasaException
+from ...exceptions import DeviceError, KasaException
 from ...interfaces.energy import Energy as EnergyInterface
 from ..smartmodule import SmartModule, raise_if_update_error
 
@@ -15,12 +15,39 @@ class Energy(SmartModule, EnergyInterface):
 
     REQUIRED_COMPONENT = "energy_monitoring"
 
+    _energy: dict[str, Any]
+    _current_consumption: float | None
+
     async def _post_update_hook(self) -> None:
-        if "voltage_mv" in self.data.get("get_emeter_data", {}):
+        try:
+            data = self.data
+        except DeviceError as de:
+            self._energy = {}
+            self._current_consumption = None
+            raise de
+
+        # If version is 1 then data is get_energy_usage
+        self._energy = data.get("get_energy_usage", data)
+
+        if "voltage_mv" in data.get("get_emeter_data", {}):
             self._supported = (
                 self._supported | EnergyInterface.ModuleFeature.VOLTAGE_CURRENT
             )
 
+        if (power := self._energy.get("current_power")) is not None or (
+            power := data.get("get_emeter_data", {}).get("power_mw")
+        ) is not None:
+            self._current_consumption = power / 1_000
+        # Fallback if get_energy_usage does not provide current_power,
+        # which can happen on some newer devices (e.g. P304M).
+        # This may not be valid scenario as it pre-dates trying get_emeter_data
+        elif (
+            power := self.data.get("get_current_power", {}).get("current_power")
+        ) is not None:
+            self._current_consumption = power
+        else:
+            self._current_consumption = None
+
     def query(self) -> dict:
         """Query to execute during the update cycle."""
         req = {
@@ -33,28 +60,21 @@ 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 or (
-            power := self.data.get("get_emeter_data", {}).get("power_mw")
-        ) is not None:
-            return power / 1_000
-        # Fallback if get_energy_usage does not provide current_power,
-        # which can happen on some newer devices (e.g. P304M).
-        elif (
-            power := self.data.get("get_current_power", {}).get("current_power")
-        ) is not None:
-            return power
-        return None
+    def optional_response_keys(self) -> list[str]:
+        """Return optional response keys for the module."""
+        if self.supported_version > 1:
+            return ["get_energy_usage"]
+        return []
+
+    @property
+    def current_consumption(self) -> float | None:
+        """Current power in watts."""
+        return self._current_consumption
 
     @property
-    @raise_if_update_error
     def energy(self) -> dict:
         """Return get_energy_usage results."""
-        if en := self.data.get("get_energy_usage"):
-            return en
-        return self.data
+        return self._energy
 
     def _get_status_from_energy(self, energy: dict) -> EmeterStatus:
         return EmeterStatus(
@@ -83,16 +103,18 @@ 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", 0) / 1_000
+        if (month := self.energy.get("month_energy")) is not None:
+            return month / 1_000
+        return None
 
     @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", 0) / 1_000
+        if (today := self.energy.get("today_energy")) is not None:
+            return today / 1_000
+        return None
 
     @property
     @raise_if_update_error

kasa/smart/smartmodule.py

@@ -72,6 +72,7 @@ class SmartModule(Module):
         self._last_update_time: float | None = None
         self._last_update_error: KasaException | None = None
         self._error_count = 0
+        self._logged_remove_keys: list[str] = []
 
     def __init_subclass__(cls, **kwargs) -> None:
         # We only want to register submodules in a modules package so that
@@ -149,6 +150,15 @@ class SmartModule(Module):
         """
         return await self._device._query_helper(method, params)
 
+    @property
+    def optional_response_keys(self) -> list[str]:
+        """Return optional response keys for the module.
+
+        Defaults to no keys. Overriding this and providing keys will remove
+        instead of raise on error.
+        """
+        return []
+
     @property
     def data(self) -> dict[str, Any]:
         """Return response data for the module.
@@ -181,12 +191,31 @@ class SmartModule(Module):
 
         filtered_data = {k: v for k, v in dev._last_update.items() if k in q_keys}
 
+        remove_keys: list[str] = []
         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 data_item in self.optional_response_keys:
+                    remove_keys.append(data_item)
+                else:
+                    raise DeviceError(
+                        f"{data_item} for {self.name}",
+                        error_code=filtered_data[data_item],
+                    )
+
+        for key in remove_keys:
+            if key not in self._logged_remove_keys:
+                self._logged_remove_keys.append(key)
+                _LOGGER.debug(
+                    "Removed key %s from response for device %s as it returned "
+                    "error: %s. This message will only be logged once per key.",
+                    key,
+                    self._device.host,
+                    filtered_data[key],
                 )
-        if len(filtered_data) == 1:
+
+            filtered_data.pop(key)
+
+        if len(filtered_data) == 1 and not remove_keys:
             return next(iter(filtered_data.values()))
 
         return filtered_data