Use decorator for AsusWrt api calls (#103690)

homeassistant/components/asuswrt/bridge.py

@@ -3,8 +3,10 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from collections import namedtuple
+from collections.abc import Awaitable, Callable, Coroutine
+import functools
 import logging
-from typing import Any, cast
+from typing import Any, TypeVar, cast
 
 from aioasuswrt.asuswrt import AsusWrt as AsusWrtLegacy
 
@@ -47,9 +49,38 @@ WrtDevice = namedtuple("WrtDevice", ["ip", "name", "connected_to"])
 _LOGGER = logging.getLogger(__name__)
 
 
-def _get_dict(keys: list, values: list) -> dict[str, Any]:
+_AsusWrtBridgeT = TypeVar("_AsusWrtBridgeT", bound="AsusWrtBridge")
-    """Create a dict from a list of keys and values."""
+_FuncType = Callable[[_AsusWrtBridgeT], Awaitable[list[Any] | dict[str, Any]]]
-    return dict(zip(keys, values))
+_ReturnFuncType = Callable[[_AsusWrtBridgeT], Coroutine[Any, Any, dict[str, Any]]]
+
+
+def handle_errors_and_zip(
+    exceptions: type[Exception] | tuple[type[Exception], ...], keys: list[str] | None
+) -> Callable[[_FuncType], _ReturnFuncType]:
+    """Run library methods and zip results or manage exceptions."""
+
+    def _handle_errors_and_zip(func: _FuncType) -> _ReturnFuncType:
+        """Run library methods and zip results or manage exceptions."""
+
+        @functools.wraps(func)
+        async def _wrapper(self: _AsusWrtBridgeT) -> dict[str, Any]:
+            try:
+                data = await func(self)
+            except exceptions as exc:
+                raise UpdateFailed(exc) from exc
+
+            if keys is None:
+                if not isinstance(data, dict):
+                    raise UpdateFailed("Received invalid data type")
+                return data
+
+            if not isinstance(data, list):
+                raise UpdateFailed("Received invalid data type")
+            return dict(zip(keys, data))
+
+        return _wrapper
+
+    return _handle_errors_and_zip
 
 
 class AsusWrtBridge(ABC):
@@ -236,38 +267,22 @@ class AsusWrtLegacyBridge(AsusWrtBridge):
         availability = await self._api.async_find_temperature_commands()
         return [SENSORS_TEMPERATURES[i] for i in range(3) if availability[i]]
 
-    async def _get_bytes(self) -> dict[str, Any]:
+    @handle_errors_and_zip((IndexError, OSError, ValueError), SENSORS_BYTES)
+    async def _get_bytes(self) -> Any:
         """Fetch byte information from the router."""
-        try:
+        return await self._api.async_get_bytes_total()
-            datas = await self._api.async_get_bytes_total()
-        except (IndexError, OSError, ValueError) as exc:
-            raise UpdateFailed(exc) from exc
 
-        return _get_dict(SENSORS_BYTES, datas)
+    @handle_errors_and_zip((IndexError, OSError, ValueError), SENSORS_RATES)
-
+    async def _get_rates(self) -> Any:
-    async def _get_rates(self) -> dict[str, Any]:
         """Fetch rates information from the router."""
-        try:
+        return await self._api.async_get_current_transfer_rates()
-            rates = await self._api.async_get_current_transfer_rates()
-        except (IndexError, OSError, ValueError) as exc:
-            raise UpdateFailed(exc) from exc
 
-        return _get_dict(SENSORS_RATES, rates)
+    @handle_errors_and_zip((IndexError, OSError, ValueError), SENSORS_LOAD_AVG)
-
+    async def _get_load_avg(self) -> Any:
-    async def _get_load_avg(self) -> dict[str, Any]:
         """Fetch load average information from the router."""
-        try:
+        return await self._api.async_get_loadavg()
-            avg = await self._api.async_get_loadavg()
-        except (IndexError, OSError, ValueError) as exc:
-            raise UpdateFailed(exc) from exc
 
-        return _get_dict(SENSORS_LOAD_AVG, avg)
+    @handle_errors_and_zip((OSError, ValueError), None)
-
+    async def _get_temperatures(self) -> Any:
-    async def _get_temperatures(self) -> dict[str, Any]:
         """Fetch temperatures information from the router."""
-        try:
+        return await self._api.async_get_temperature()
-            temperatures: dict[str, Any] = await self._api.async_get_temperature()
-        except (OSError, ValueError) as exc:
-            raise UpdateFailed(exc) from exc
-
-        return temperatures

tests/components/asuswrt/conftest.py

@@ -13,7 +13,7 @@ ASUSWRT_LEGACY_LIB = f"{ASUSWRT_BASE}.bridge.AsusWrtLegacy"
 MOCK_BYTES_TOTAL = [60000000000, 50000000000]
 MOCK_CURRENT_TRANSFER_RATES = [20000000, 10000000]
 MOCK_LOAD_AVG = [1.1, 1.2, 1.3]
-MOCK_TEMPERATURES = {"2.4GHz": 40.2, "CPU": 71.2}
+MOCK_TEMPERATURES = {"2.4GHz": 40.2, "5.0GHz": 0, "CPU": 71.2}
 
 
 @pytest.fixture(name="patch_setup_entry")

tests/components/asuswrt/test_sensor.py

@@ -302,3 +302,28 @@ async def test_unique_id_migration(
     migr_entity = entity_registry.async_get(f"{sensor.DOMAIN}.{obj_entity_id}")
     assert migr_entity is not None
     assert migr_entity.unique_id == slugify(f"{ROUTER_MAC_ADDR}_sensor_tx_bytes")
+
+
+async def test_decorator_errors(
+    hass: HomeAssistant, connect_legacy, mock_available_temps
+) -> None:
+    """Test AsusWRT sensors are unavailable on decorator type check error."""
+    sensors = [*SENSORS_BYTES, *SENSORS_TEMPERATURES]
+    config_entry, sensor_prefix = _setup_entry(hass, CONFIG_DATA_TELNET, sensors)
+    config_entry.add_to_hass(hass)
+
+    mock_available_temps[1] = True
+    connect_legacy.return_value.async_get_bytes_total.return_value = "bad_response"
+    connect_legacy.return_value.async_get_temperature.return_value = "bad_response"
+
+    # initial devices setup
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+    async_fire_time_changed(hass, utcnow() + timedelta(seconds=30))
+    await hass.async_block_till_done()
+
+    for sensor_name in sensors:
+        assert (
+            hass.states.get(f"{sensor_prefix}_{slugify(sensor_name)}").state
+            == STATE_UNAVAILABLE
+        )