From 03aeaba7ef6fc0093cca1d9e9984796b0d473ad4 Mon Sep 17 00:00:00 2001 From: Michael <35783820+mib1185@users.noreply.github.com> Date: Wed, 22 Mar 2023 22:34:23 +0100 Subject: [PATCH] Turn AVM FRITZ!Box Tools sensors into coordinator entities (#89953) * make sensors coordinator entities * apply suggestions * move _attr_has_entity_name up --- homeassistant/components/fritz/common.py | 85 +++++++++++++++++++++++- homeassistant/components/fritz/sensor.py | 53 ++++----------- 2 files changed, 95 insertions(+), 43 deletions(-) diff --git a/homeassistant/components/fritz/common.py b/homeassistant/components/fritz/common.py index 09103a0bcc8..f6025e773e0 100644 --- a/homeassistant/components/fritz/common.py +++ b/homeassistant/components/fritz/common.py @@ -35,7 +35,8 @@ from homeassistant.helpers import ( update_coordinator, ) from homeassistant.helpers.dispatcher import async_dispatcher_send -from homeassistant.helpers.entity import DeviceInfo +from homeassistant.helpers.entity import DeviceInfo, EntityDescription +from homeassistant.helpers.typing import StateType from homeassistant.util import dt as dt_util from .const import ( @@ -136,7 +137,9 @@ class HostInfo(TypedDict): status: bool -class FritzBoxTools(update_coordinator.DataUpdateCoordinator[None]): +class FritzBoxTools( + update_coordinator.DataUpdateCoordinator[dict[str, bool | StateType]] +): """FritzBoxTools class.""" def __init__( @@ -175,6 +178,9 @@ class FritzBoxTools(update_coordinator.DataUpdateCoordinator[None]): self._latest_firmware: str | None = None self._update_available: bool = False self._release_url: str | None = None + self._entity_update_functions: dict[ + str, Callable[[FritzStatus, StateType], Any] + ] = {} async def async_setup( self, options: MappingProxyType[str, Any] | None = None @@ -237,12 +243,36 @@ class FritzBoxTools(update_coordinator.DataUpdateCoordinator[None]): ) self.device_is_router = self.fritz_status.has_wan_enabled - async def _async_update_data(self) -> None: + def register_entity_updates( + self, key: str, update_fn: Callable[[FritzStatus, StateType], Any] + ) -> Callable[[], None]: + """Register an entity to be updated by coordinator.""" + + def unregister_entity_updates() -> None: + """Unregister an entity to be updated by coordinator.""" + if key in self._entity_update_functions: + _LOGGER.debug("unregister entity %s from updates", key) + self._entity_update_functions.pop(key) + + if key not in self._entity_update_functions: + _LOGGER.debug("register entity %s for updates", key) + self._entity_update_functions[key] = update_fn + return unregister_entity_updates + + async def _async_update_data(self) -> dict[str, bool | StateType]: """Update FritzboxTools data.""" + enity_data: dict[str, bool | StateType] = {} try: await self.async_scan_devices() + for key, update_fn in self._entity_update_functions.items(): + _LOGGER.debug("update entity %s", key) + enity_data[key] = await self.hass.async_add_executor_job( + update_fn, self.fritz_status, self.data.get(key) + ) except FRITZ_EXCEPTIONS as ex: raise update_coordinator.UpdateFailed(ex) from ex + _LOGGER.debug("enity_data: %s", enity_data) + return enity_data @property def unique_id(self) -> str: @@ -981,6 +1011,55 @@ class FritzBoxBaseEntity: ) +@dataclass +class FritzRequireKeysMixin: + """Fritz entity description mix in.""" + + value_fn: Callable[[FritzStatus, Any], Any] + + +@dataclass +class FritzEntityDescription(EntityDescription, FritzRequireKeysMixin): + """Fritz entity base description.""" + + +class FritzBoxBaseCoordinatorEntity(update_coordinator.CoordinatorEntity): + """Fritz host coordinator entity base class.""" + + coordinator: AvmWrapper + entity_description: FritzEntityDescription + _attr_has_entity_name = True + + def __init__( + self, + avm_wrapper: AvmWrapper, + device_name: str, + description: FritzEntityDescription, + ) -> None: + """Init device info class.""" + super().__init__(avm_wrapper) + self.async_on_remove( + avm_wrapper.register_entity_updates(description.key, description.value_fn) + ) + self.entity_description = description + self._device_name = device_name + self._attr_name = description.name + self._attr_unique_id = f"{avm_wrapper.unique_id}-{description.key}" + + @property + def device_info(self) -> DeviceInfo: + """Return the device information.""" + return DeviceInfo( + configuration_url=f"http://{self.coordinator.host}", + connections={(dr.CONNECTION_NETWORK_MAC, self.coordinator.mac)}, + identifiers={(DOMAIN, self.coordinator.unique_id)}, + manufacturer="AVM", + model=self.coordinator.model, + name=self._device_name, + sw_version=self.coordinator.current_firmware, + ) + + @dataclass class ConnectionInfo: """Fritz sensor connection information class.""" diff --git a/homeassistant/components/fritz/sensor.py b/homeassistant/components/fritz/sensor.py index 628d56dc450..4b15f3f92de 100644 --- a/homeassistant/components/fritz/sensor.py +++ b/homeassistant/components/fritz/sensor.py @@ -5,9 +5,7 @@ from collections.abc import Callable from dataclasses import dataclass from datetime import datetime, timedelta import logging -from typing import Any -from fritzconnection.core.exceptions import FritzConnectionException from fritzconnection.lib.fritzstatus import FritzStatus from homeassistant.components.sensor import ( @@ -25,9 +23,15 @@ from homeassistant.const import ( ) from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback +from homeassistant.helpers.typing import StateType from homeassistant.util.dt import utcnow -from .common import AvmWrapper, ConnectionInfo, FritzBoxBaseEntity +from .common import ( + AvmWrapper, + ConnectionInfo, + FritzBoxBaseCoordinatorEntity, + FritzEntityDescription, +) from .const import DOMAIN, DSL_CONNECTION, UPTIME_DEVIATION _LOGGER = logging.getLogger(__name__) @@ -139,14 +143,7 @@ def _retrieve_link_attenuation_received_state( @dataclass -class FritzRequireKeysMixin: - """Fritz sensor data class.""" - - value_fn: Callable[[FritzStatus, Any], Any] - - -@dataclass -class FritzSensorEntityDescription(SensorEntityDescription, FritzRequireKeysMixin): +class FritzSensorEntityDescription(SensorEntityDescription, FritzEntityDescription): """Describes Fritz sensor entity.""" is_suitable: Callable[[ConnectionInfo], bool] = lambda info: info.wan_enabled @@ -304,36 +301,12 @@ async def async_setup_entry( async_add_entities(entities, True) -class FritzBoxSensor(FritzBoxBaseEntity, SensorEntity): +class FritzBoxSensor(FritzBoxBaseCoordinatorEntity, SensorEntity): """Define FRITZ!Box connectivity class.""" entity_description: FritzSensorEntityDescription - def __init__( - self, - avm_wrapper: AvmWrapper, - device_friendly_name: str, - description: FritzSensorEntityDescription, - ) -> None: - """Init FRITZ!Box connectivity class.""" - self.entity_description = description - self._last_device_value: str | None = None - self._attr_available = True - self._attr_name = f"{device_friendly_name} {description.name}" - self._attr_unique_id = f"{avm_wrapper.unique_id}-{description.key}" - super().__init__(avm_wrapper, device_friendly_name) - - def update(self) -> None: - """Update data.""" - _LOGGER.debug("Updating FRITZ!Box sensors") - - status: FritzStatus = self._avm_wrapper.fritz_status - try: - self._attr_native_value = ( - self._last_device_value - ) = self.entity_description.value_fn(status, self._last_device_value) - except FritzConnectionException: - _LOGGER.error("Error getting the state from the FRITZ!Box", exc_info=True) - self._attr_available = False - return - self._attr_available = True + @property + def native_value(self) -> StateType: + """Return the value reported by the sensor.""" + return self.coordinator.data.get(self.entity_description.key)