From fe88b52e19534ad84e5595070136bcf1e1adb0be Mon Sep 17 00:00:00 2001 From: "Teemu R." Date: Fri, 20 Dec 2024 09:53:07 +0100 Subject: [PATCH] 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 | 72 +++++++++++++++--------- kasa/smart/smartmodule.py | 35 +++++++++++- tests/smart/modules/test_energy.py | 90 +++++++++++++++++++++++++++++- tests/smart/test_smartdevice.py | 4 +- 4 files changed, 170 insertions(+), 31 deletions(-) diff --git a/kasa/smart/modules/energy.py b/kasa/smart/modules/energy.py index 6b5bdb57..0cfdc92c 100644 --- a/kasa/smart/modules/energy.py +++ b/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 diff --git a/kasa/smart/smartmodule.py b/kasa/smart/smartmodule.py index 31fc8f35..a5666f63 100644 --- a/kasa/smart/smartmodule.py +++ b/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 diff --git a/tests/smart/modules/test_energy.py b/tests/smart/modules/test_energy.py index fdbea88b..7b31d74b 100644 --- a/tests/smart/modules/test_energy.py +++ b/tests/smart/modules/test_energy.py @@ -1,7 +1,14 @@ +import copy +import logging +from contextlib import nullcontext as does_not_raise +from unittest.mock import patch + import pytest -from kasa import Module, SmartDevice +from kasa import DeviceError, Module +from kasa.exceptions import SmartErrorCode from kasa.interfaces.energy import Energy +from kasa.smart import SmartDevice from kasa.smart.modules import Energy as SmartEnergyModule from tests.conftest import has_emeter_smart @@ -19,3 +26,84 @@ async def test_supported(dev: SmartDevice): assert energy_module.supports(Energy.ModuleFeature.VOLTAGE_CURRENT) is False else: assert energy_module.supports(Energy.ModuleFeature.VOLTAGE_CURRENT) is True + + +@has_emeter_smart +async def test_get_energy_usage_error( + dev: SmartDevice, caplog: pytest.LogCaptureFixture +): + """Test errors on get_energy_usage.""" + caplog.set_level(logging.DEBUG) + + energy_module = dev.modules.get(Module.Energy) + if not energy_module: + pytest.skip(f"Energy module not supported for {dev}.") + + version = dev._components["energy_monitoring"] + + expected_raise = does_not_raise() if version > 1 else pytest.raises(DeviceError) + if version > 1: + expected = "get_energy_usage" + expected_current_consumption = 2.002 + else: + expected = "current_power" + expected_current_consumption = None + + assert expected in energy_module.data + assert energy_module.current_consumption is not None + assert energy_module.consumption_today is not None + assert energy_module.consumption_this_month is not None + + last_update = copy.deepcopy(dev._last_update) + resp = copy.deepcopy(last_update) + + if ed := resp.get("get_emeter_data"): + ed["power_mw"] = 2002 + if cp := resp.get("get_current_power"): + cp["current_power"] = 2.002 + resp["get_energy_usage"] = SmartErrorCode.JSON_DECODE_FAIL_ERROR + + # version 1 only has get_energy_usage so module should raise an error if + # version 1 and get_energy_usage is in error + with patch.object(dev.protocol, "query", return_value=resp): + await dev.update() + + with expected_raise: + assert "get_energy_usage" not in energy_module.data + + assert energy_module.current_consumption == expected_current_consumption + assert energy_module.consumption_today is None + assert energy_module.consumption_this_month is None + + msg = ( + f"Removed key get_energy_usage from response for device {dev.host}" + " as it returned error: JSON_DECODE_FAIL_ERROR" + ) + if version > 1: + assert msg in caplog.text + + # Now test with no get_emeter_data + # This may not be valid scenario but we have a fallback to get_current_power + # just in case that should be tested. + caplog.clear() + resp = copy.deepcopy(last_update) + + if cp := resp.get("get_current_power"): + cp["current_power"] = 2.002 + resp["get_energy_usage"] = SmartErrorCode.JSON_DECODE_FAIL_ERROR + + # Remove get_emeter_data from the response and from the device which will + # remember it otherwise. + resp.pop("get_emeter_data", None) + dev._last_update.pop("get_emeter_data", None) + + with patch.object(dev.protocol, "query", return_value=resp): + await dev.update() + + with expected_raise: + assert "get_energy_usage" not in energy_module.data + + assert energy_module.current_consumption == expected_current_consumption + + # message should only be logged once + assert msg not in caplog.text diff --git a/tests/smart/test_smartdevice.py b/tests/smart/test_smartdevice.py index 83635d8e..549eb8ad 100644 --- a/tests/smart/test_smartdevice.py +++ b/tests/smart/test_smartdevice.py @@ -355,7 +355,7 @@ async def test_update_module_query_errors( if mod.name == "Energy": emod = cast(Energy, mod) with pytest.raises(KasaException, match="Module update error"): - assert emod.current_consumption is not None + assert emod.status is not None else: assert mod.disabled is False assert mod._error_count == 0 @@ -363,7 +363,7 @@ async def test_update_module_query_errors( # Test one of the raise_if_update_error doesn't raise if mod.name == "Energy": emod = cast(Energy, mod) - assert emod.current_consumption is not None + assert emod.status is not None async def test_get_modules():