Implement DataUpdateCoordinator for Fritz (#60909)

* Implement DataUpdateCoordinator for Fritz

* mypy

* Wrap sync method to async

* Apply review comments + final cleanup

* CoordinatorEntity

homeassistant/components/fritz/__init__.py

@@ -5,14 +5,8 @@ from fritzconnection.core.exceptions import FritzConnectionException, FritzSecur
 from fritzconnection.core.logger import fritzlogger
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import (
-    CONF_HOST,
-    CONF_PASSWORD,
-    CONF_PORT,
-    CONF_USERNAME,
-    EVENT_HOMEASSISTANT_STOP,
-)
-from homeassistant.core import Event, HomeAssistant, callback
+from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME
+from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 
 from .common import FritzBoxTools, FritzData
@@ -41,8 +35,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
 
     try:
-        await fritz_tools.async_setup()
-        await fritz_tools.async_start(entry.options)
+        await fritz_tools.async_setup(entry.options)
     except FritzSecurityError as ex:
         raise ConfigEntryAuthFailed from ex
     except FritzConnectionException as ex:
@@ -54,15 +47,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     if DATA_FRITZ not in hass.data:
         hass.data[DATA_FRITZ] = FritzData()
 
-    @callback
-    def _async_unload(event: Event) -> None:
-        fritz_tools.async_unload()
-
-    entry.async_on_unload(
-        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _async_unload)
-    )
     entry.async_on_unload(entry.add_update_listener(update_listener))
 
+    await fritz_tools.async_config_entry_first_refresh()
+
     # Load the other platforms like switch
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
@@ -74,7 +62,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload FRITZ!Box Tools config entry."""
     fritzbox: FritzBoxTools = hass.data[DOMAIN][entry.entry_id]
-    fritzbox.async_unload()
 
     fritz_data = hass.data[DATA_FRITZ]
     fritz_data.tracked.pop(fritzbox.unique_id)

homeassistant/components/fritz/common.py

@@ -12,6 +12,7 @@ from fritzconnection import FritzConnection
 from fritzconnection.core.exceptions import (
     FritzActionError,
     FritzConnectionException,
+    FritzSecurityError,
     FritzServiceError,
 )
 from fritzconnection.lib.fritzhosts import FritzHosts
@@ -24,21 +25,21 @@ from homeassistant.components.device_tracker.const import (
 )
 from homeassistant.components.switch import DOMAIN as DEVICE_SWITCH_DOMAIN
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.core import CALLBACK_TYPE, HomeAssistant, ServiceCall, callback
+from homeassistant.core import HomeAssistant, ServiceCall, callback
 from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers import update_coordinator
 from homeassistant.helpers.device_registry import (
     CONNECTION_NETWORK_MAC,
     async_entries_for_config_entry,
     async_get,
 )
-from homeassistant.helpers.dispatcher import async_dispatcher_connect, dispatcher_send
-from homeassistant.helpers.entity import DeviceInfo, Entity
+from homeassistant.helpers.dispatcher import dispatcher_send
+from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_registry import (
     EntityRegistry,
     RegistryEntry,
     async_entries_for_device,
 )
-from homeassistant.helpers.event import async_track_time_interval
 from homeassistant.util import dt as dt_util
 
 from .const import (
@@ -50,7 +51,6 @@ from .const import (
     SERVICE_CLEANUP,
     SERVICE_REBOOT,
     SERVICE_RECONNECT,
-    TRACKER_SCAN_INTERVAL,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -105,6 +105,7 @@ class Device:
     mac: str
     ip_address: str
     name: str
+    wan_access: bool
 
 
 class HostInfo(TypedDict):
@@ -116,7 +117,7 @@ class HostInfo(TypedDict):
     status: bool
 
 
-class FritzBoxTools:
+class FritzBoxTools(update_coordinator.DataUpdateCoordinator):
     """FrtizBoxTools class."""
 
     def __init__(
@@ -128,7 +129,13 @@ class FritzBoxTools:
         port: int = DEFAULT_PORT,
     ) -> None:
         """Initialize FritzboxTools class."""
-        self._cancel_scan: CALLBACK_TYPE | None = None
+        super().__init__(
+            hass=hass,
+            logger=_LOGGER,
+            name=f"{DOMAIN}-{host}-coordinator",
+            update_interval=timedelta(seconds=30),
+        )
+
         self._devices: dict[str, FritzDevice] = {}
         self._options: MappingProxyType[str, Any] | None = None
         self._unique_id: str | None = None
@@ -146,8 +153,11 @@ class FritzBoxTools:
         self._latest_firmware: str | None = None
         self._update_available: bool = False
 
-    async def async_setup(self) -> None:
+    async def async_setup(
+        self, options: MappingProxyType[str, Any] | None = None
+    ) -> None:
         """Wrap up FritzboxTools class setup."""
+        self._options = options
         await self.hass.async_add_executor_job(self.setup)
 
     def setup(self) -> None:
@@ -175,23 +185,14 @@ class FritzBoxTools:
 
         self._update_available, self._latest_firmware = self._update_device_info()
 
-    async def async_start(self, options: MappingProxyType[str, Any]) -> None:
-        """Start FritzHosts connection."""
-        self.fritz_hosts = FritzHosts(fc=self.connection)
-        self._options = options
-        await self.hass.async_add_executor_job(self.scan_devices)
-
-        self._cancel_scan = async_track_time_interval(
-            self.hass, self.scan_devices, timedelta(seconds=TRACKER_SCAN_INTERVAL)
-        )
-
     @callback
-    def async_unload(self) -> None:
-        """Unload FritzboxTools class."""
-        _LOGGER.debug("Unloading FRITZ!Box router integration")
-        if self._cancel_scan is not None:
-            self._cancel_scan()
-            self._cancel_scan = None
+    async def _async_update_data(self) -> None:
+        """Update FritzboxTools data."""
+        try:
+            self.fritz_hosts = FritzHosts(fc=self.connection)
+            await self.async_scan_devices()
+        except (FritzSecurityError, FritzConnectionException) as ex:
+            raise update_coordinator.UpdateFailed from ex
 
     @property
     def unique_id(self) -> str:
@@ -262,6 +263,10 @@ class FritzBoxTools:
         )
         return bool(version), version
 
+    async def async_scan_devices(self, now: datetime | None = None) -> None:
+        """Wrap up FritzboxTools class scan."""
+        await self.hass.async_add_executor_job(self.scan_devices, now)
+
     def scan_devices(self, now: datetime | None = None) -> None:
         """Scan for new devices and return a list of found device ids."""
         _LOGGER.debug("Checking devices for FRITZ!Box router %s", self.host)
@@ -283,8 +288,17 @@ class FritzBoxTools:
             dev_name = known_host["name"]
             dev_ip = known_host["ip"]
             dev_home = known_host["status"]
+            dev_wan_access = True
+            if dev_ip:
+                wan_access = self.connection.call_action(
+                    "X_AVM-DE_HostFilter:1",
+                    "GetWANAccessByIP",
+                    NewIPv4Address=dev_ip,
+                )
+                if wan_access:
+                    dev_wan_access = not wan_access.get("NewDisallow")
 
-            dev_info = Device(dev_mac, dev_ip, dev_name)
+            dev_info = Device(dev_mac, dev_ip, dev_name, dev_wan_access)
 
             if dev_mac in self._devices:
                 self._devices[dev_mac].update(dev_info, dev_home, consider_home)
@@ -387,11 +401,12 @@ class FritzData:
     profile_switches: dict = field(default_factory=dict)
 
 
-class FritzDeviceBase(Entity):
+class FritzDeviceBase(update_coordinator.CoordinatorEntity):
     """Entity base class for a device connected to a FRITZ!Box router."""
 
     def __init__(self, router: FritzBoxTools, device: FritzDevice) -> None:
         """Initialize a FRITZ!Box device."""
+        super().__init__(router)
         self._router = router
         self._mac: str = device.mac_address
         self._name: str = device.hostname or DEFAULT_DEVICE_NAME
@@ -405,8 +420,7 @@ class FritzDeviceBase(Entity):
     def ip_address(self) -> str | None:
         """Return the primary ip address of the device."""
         if self._mac:
-            device: FritzDevice = self._router.devices[self._mac]
-            return device.ip_address
+            return self._router.devices[self._mac].ip_address
         return None
 
     @property
@@ -418,8 +432,7 @@ class FritzDeviceBase(Entity):
     def hostname(self) -> str | None:
         """Return hostname of the device."""
         if self._mac:
-            device: FritzDevice = self._router.devices[self._mac]
-            return device.hostname
+            return self._router.devices[self._mac].hostname
         return None
 
     @property
@@ -451,17 +464,6 @@ class FritzDeviceBase(Entity):
         await self.async_process_update()
         self.async_write_ha_state()
 
-    async def async_added_to_hass(self) -> None:
-        """Register state update callback."""
-        await self.async_process_update()
-        self.async_on_remove(
-            async_dispatcher_connect(
-                self.hass,
-                self._router.signal_device_update,
-                self.async_on_demand_update,
-            )
-        )
-
 
 class FritzDevice:
     """Representation of a device connected to the FRITZ!Box."""
@@ -473,6 +475,7 @@ class FritzDevice:
         self._ip_address: str | None = None
         self._last_activity: datetime | None = None
         self._connected = False
+        self._wan_access = False
 
     def update(self, dev_info: Device, dev_home: bool, consider_home: float) -> None:
         """Update device info."""
@@ -494,6 +497,7 @@ class FritzDevice:
             self._last_activity = utc_point_in_time
 
         self._ip_address = dev_info.ip_address
+        self._wan_access = dev_info.wan_access
 
     @property
     def is_connected(self) -> bool:
@@ -520,6 +524,11 @@ class FritzDevice:
         """Return device last activity."""
         return self._last_activity
 
+    @property
+    def wan_access(self) -> bool:
+        """Return device wan access."""
+        return self._wan_access
+
 
 class SwitchInfo(TypedDict):
     """FRITZ!Box switch info class."""

homeassistant/components/fritz/const.py

@@ -35,6 +35,4 @@ SWITCH_TYPE_DEFLECTION = "CallDeflection"
 SWITCH_TYPE_PORTFORWARD = "PortForward"
 SWITCH_TYPE_WIFINETWORK = "WiFiNetwork"
 
-TRACKER_SCAN_INTERVAL = 30
-
 UPTIME_DEVIATION = 5

homeassistant/components/fritz/device_tracker.py

@@ -119,12 +119,11 @@ class FritzBoxTracker(FritzDeviceBase, ScannerEntity):
         """Initialize a FRITZ!Box device."""
         super().__init__(router, device)
         self._last_activity: datetime.datetime | None = device.last_activity
-        self._active = False
 
     @property
     def is_connected(self) -> bool:
         """Return device status."""
-        return self._active
+        return self._router.devices[self._mac].is_connected
 
     @property
     def unique_id(self) -> str:
@@ -142,6 +141,7 @@ class FritzBoxTracker(FritzDeviceBase, ScannerEntity):
     def extra_state_attributes(self) -> dict[str, str]:
         """Return the attributes."""
         attrs: dict[str, str] = {}
+        self._last_activity = self._router.devices[self._mac].last_activity
         if self._last_activity is not None:
             attrs["last_time_reachable"] = self._last_activity.isoformat(
                 timespec="seconds"
@@ -152,12 +152,3 @@ class FritzBoxTracker(FritzDeviceBase, ScannerEntity):
     def source_type(self) -> str:
         """Return tracker source type."""
         return SOURCE_TYPE_ROUTER
-
-    async def async_process_update(self) -> None:
-        """Update device."""
-        if not self._mac:
-            return
-
-        device = self._router.devices[self._mac]
-        self._active = device.is_connected
-        self._last_activity = device.last_activity

homeassistant/components/fritz/switch.py

@@ -595,23 +595,10 @@ class FritzBoxProfileSwitch(FritzDeviceBase, SwitchEntity):
         self._attr_unique_id = f"{self._mac}_internet_access"
         self._attr_entity_category = EntityCategory.CONFIG
 
-    async def async_process_update(self) -> None:
-        """Update device."""
-        if not self._mac or not self.ip_address:
-            return
-
-        wan_disable_info = await async_service_call_action(
-            self._router,
-            "X_AVM-DE_HostFilter",
-            "1",
-            "GetWANAccessByIP",
-            NewIPv4Address=self.ip_address,
-        )
-
-        if wan_disable_info is None:
-            return
-
-        self._attr_is_on = not wan_disable_info["NewDisallow"]
+    @property
+    def is_on(self) -> bool:
+        """Switch status."""
+        return self._router.devices[self._mac].wan_access
 
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn on switch."""
@@ -624,7 +611,6 @@ class FritzBoxProfileSwitch(FritzDeviceBase, SwitchEntity):
     async def _async_handle_turn_on_off(self, turn_on: bool) -> bool:
         """Handle switch state change request."""
         await self._async_switch_on_off(turn_on)
-        self._attr_is_on = turn_on
         self.async_write_ha_state()
         return True