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

@@ -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."""

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)