Add base NetgearDataCoordinator to netgear (#164816)

homeassistant/components/netgear/__init__.py

@@ -10,10 +10,14 @@ from homeassistant.const import CONF_PORT, CONF_SSL
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import device_registry as dr, entity_registry as er
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from .const import PLATFORMS
-from .coordinator import NetgearConfigEntry, NetgearRuntimeData
+from .coordinator import (
+    NetgearConfigEntry,
+    NetgearDataCoordinator,
+    NetgearFirmwareCoordinator,
+    NetgearRuntimeData,
+)
 from .errors import CannotLoginException
 from .router import NetgearRouter
 
@@ -21,7 +25,6 @@ _LOGGER = logging.getLogger(__name__)
 
 SCAN_INTERVAL = timedelta(seconds=30)
 SPEED_TEST_INTERVAL = timedelta(hours=2)
-SCAN_INTERVAL_FIRMWARE = timedelta(hours=5)
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: NetgearConfigEntry) -> bool:
@@ -63,10 +66,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: NetgearConfigEntry) -> b
         """Fetch data from the router."""
         return await router.async_get_speed_test()
 
-    async def async_check_firmware() -> dict[str, Any] | None:
-        """Check for new firmware of the router."""
-        return await router.async_check_new_firmware()
-
     async def async_update_utilization() -> dict[str, Any] | None:
         """Fetch data from the router."""
         return await router.async_get_utilization()
@@ -76,57 +75,50 @@ async def async_setup_entry(hass: HomeAssistant, entry: NetgearConfigEntry) -> b
         return await router.async_get_link_status()
 
     # Create update coordinators
-    coordinator = DataUpdateCoordinator(
+    coordinator_tracker = NetgearDataCoordinator(
         hass,
-        _LOGGER,
-        config_entry=entry,
-        name=f"{router.device_name} Devices",
+        router,
+        entry,
+        name="Devices",
         update_method=async_update_devices,
         update_interval=SCAN_INTERVAL,
     )
-    coordinator_traffic_meter = DataUpdateCoordinator(
+    coordinator_traffic_meter = NetgearDataCoordinator(
         hass,
-        _LOGGER,
-        config_entry=entry,
-        name=f"{router.device_name} Traffic meter",
+        router,
+        entry,
+        name="Traffic meter",
         update_method=async_update_traffic_meter,
         update_interval=SCAN_INTERVAL,
     )
-    coordinator_speed_test = DataUpdateCoordinator(
+    coordinator_speed_test = NetgearDataCoordinator(
         hass,
-        _LOGGER,
-        config_entry=entry,
-        name=f"{router.device_name} Speed test",
+        router,
+        entry,
+        name="Speed test",
         update_method=async_update_speed_test,
         update_interval=SPEED_TEST_INTERVAL,
     )
-    coordinator_firmware = DataUpdateCoordinator(
+    coordinator_firmware = NetgearFirmwareCoordinator(hass, router, entry)
+    coordinator_utilization = NetgearDataCoordinator(
         hass,
-        _LOGGER,
-        config_entry=entry,
-        name=f"{router.device_name} Firmware",
-        update_method=async_check_firmware,
-        update_interval=SCAN_INTERVAL_FIRMWARE,
-    )
-    coordinator_utilization = DataUpdateCoordinator(
-        hass,
-        _LOGGER,
-        config_entry=entry,
-        name=f"{router.device_name} Utilization",
+        router,
+        entry,
+        name="Utilization",
         update_method=async_update_utilization,
         update_interval=SCAN_INTERVAL,
     )
-    coordinator_link = DataUpdateCoordinator(
+    coordinator_link = NetgearDataCoordinator(
         hass,
-        _LOGGER,
-        config_entry=entry,
-        name=f"{router.device_name} Ethernet Link Status",
+        router,
+        entry,
+        name="Ethernet Link Status",
         update_method=async_check_link_status,
         update_interval=SCAN_INTERVAL,
     )
 
     if router.track_devices:
-        await coordinator.async_config_entry_first_refresh()
+        await coordinator_tracker.async_config_entry_first_refresh()
     await coordinator_traffic_meter.async_config_entry_first_refresh()
     await coordinator_firmware.async_config_entry_first_refresh()
     await coordinator_utilization.async_config_entry_first_refresh()
@@ -134,7 +126,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: NetgearConfigEntry) -> b
 
     entry.runtime_data = NetgearRuntimeData(
         router=router,
-        coordinator=coordinator,
+        coordinator_tracker=coordinator_tracker,
         coordinator_traffic=coordinator_traffic_meter,
         coordinator_speed=coordinator_speed_test,
         coordinator_firmware=coordinator_firmware,

homeassistant/components/netgear/button.py

@@ -12,9 +12,8 @@ from homeassistant.components.button import (
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .coordinator import NetgearConfigEntry
+from .coordinator import NetgearConfigEntry, NetgearDataCoordinator
 from .entity import NetgearRouterCoordinatorEntity
 from .router import NetgearRouter
 
@@ -43,9 +42,9 @@ async def async_setup_entry(
 ) -> None:
     """Set up button for Netgear component."""
     router = entry.runtime_data.router
-    coordinator = entry.runtime_data.coordinator
+    coordinator_tracker = entry.runtime_data.coordinator_tracker
     async_add_entities(
-        NetgearRouterButtonEntity(coordinator, router, entity_description)
+        NetgearRouterButtonEntity(coordinator_tracker, router, entity_description)
         for entity_description in BUTTONS
     )
 
@@ -57,7 +56,7 @@ class NetgearRouterButtonEntity(NetgearRouterCoordinatorEntity, ButtonEntity):
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator[bool],
+        coordinator: NetgearDataCoordinator[bool],
         router: NetgearRouter,
         entity_description: NetgearButtonEntityDescription,
     ) -> None:

homeassistant/components/netgear/coordinator.py

@@ -2,26 +2,81 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable, Coroutine
 from dataclasses import dataclass
+from datetime import timedelta
+import logging
 from typing import Any
 
 from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from .router import NetgearRouter
 
+_LOGGER = logging.getLogger(__name__)
+
+SCAN_INTERVAL_FIRMWARE = timedelta(hours=5)
+
 
 @dataclass
 class NetgearRuntimeData:
     """Runtime data for the Netgear integration."""
 
     router: NetgearRouter
-    coordinator: DataUpdateCoordinator[bool]
-    coordinator_traffic: DataUpdateCoordinator[dict[str, Any] | None]
-    coordinator_speed: DataUpdateCoordinator[dict[str, Any] | None]
-    coordinator_firmware: DataUpdateCoordinator[dict[str, Any] | None]
-    coordinator_utilization: DataUpdateCoordinator[dict[str, Any] | None]
-    coordinator_link: DataUpdateCoordinator[dict[str, Any] | None]
+    coordinator_tracker: NetgearDataCoordinator[bool]
+    coordinator_traffic: NetgearDataCoordinator[dict[str, Any] | None]
+    coordinator_speed: NetgearDataCoordinator[dict[str, Any] | None]
+    coordinator_firmware: NetgearFirmwareCoordinator
+    coordinator_utilization: NetgearDataCoordinator[dict[str, Any] | None]
+    coordinator_link: NetgearDataCoordinator[dict[str, Any] | None]
 
 
 type NetgearConfigEntry = ConfigEntry[NetgearRuntimeData]
+
+
+class NetgearDataCoordinator[T](DataUpdateCoordinator[T]):
+    """Base coordinator for Netgear."""
+
+    config_entry: NetgearConfigEntry
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        router: NetgearRouter,
+        entry: NetgearConfigEntry,
+        *,
+        name: str,
+        update_interval: timedelta,
+        update_method: Callable[[], Coroutine[Any, Any, T]] | None = None,
+    ) -> None:
+        """Initialize the coordinator."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            config_entry=entry,
+            name=f"{router.device_name} {name}",
+            update_interval=update_interval,
+            update_method=update_method,
+        )
+        self.router = router
+
+
+class NetgearFirmwareCoordinator(NetgearDataCoordinator[dict[str, Any] | None]):
+    """Coordinator for Netgear firmware updates."""
+
+    def __init__(
+        self, hass: HomeAssistant, router: NetgearRouter, entry: NetgearConfigEntry
+    ) -> None:
+        """Initialize the coordinator."""
+        super().__init__(
+            hass,
+            router,
+            entry,
+            name="Firmware",
+            update_interval=SCAN_INTERVAL_FIRMWARE,
+        )
+
+    async def _async_update_data(self) -> dict[str, Any] | None:
+        """Check for new firmware of the router."""
+        return await self.router.async_check_new_firmware()

homeassistant/components/netgear/device_tracker.py

@@ -7,10 +7,9 @@ import logging
 from homeassistant.components.device_tracker import ScannerEntity
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from .const import DEVICE_ICONS
-from .coordinator import NetgearConfigEntry
+from .coordinator import NetgearConfigEntry, NetgearDataCoordinator
 from .entity import NetgearDeviceEntity
 from .router import NetgearRouter
 
@@ -24,13 +23,13 @@ async def async_setup_entry(
 ) -> None:
     """Set up device tracker for Netgear component."""
     router = entry.runtime_data.router
-    coordinator = entry.runtime_data.coordinator
+    coordinator_tracker = entry.runtime_data.coordinator_tracker
     tracked = set()
 
     @callback
     def new_device_callback() -> None:
         """Add new devices if needed."""
-        if not coordinator.data:
+        if not coordinator_tracker.data:
             return
 
         new_entities = []
@@ -39,14 +38,16 @@ async def async_setup_entry(
             if mac in tracked:
                 continue
 
-            new_entities.append(NetgearScannerEntity(coordinator, router, device))
+            new_entities.append(
+                NetgearScannerEntity(coordinator_tracker, router, device)
+            )
             tracked.add(mac)
 
         async_add_entities(new_entities)
 
-    entry.async_on_unload(coordinator.async_add_listener(new_device_callback))
+    entry.async_on_unload(coordinator_tracker.async_add_listener(new_device_callback))
 
-    coordinator.data = True
+    coordinator_tracker.data = True
     new_device_callback()
 
 
@@ -57,7 +58,7 @@ class NetgearScannerEntity(NetgearDeviceEntity, ScannerEntity):
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator[bool],
+        coordinator: NetgearDataCoordinator[bool],
         router: NetgearRouter,
         device: dict,
     ) -> None:

homeassistant/components/netgear/entity.py

@@ -10,12 +10,10 @@ from homeassistant.core import callback
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity import Entity
-from homeassistant.helpers.update_coordinator import (
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-)
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import DOMAIN
+from .coordinator import NetgearDataCoordinator
 from .router import NetgearRouter
 
 
@@ -26,7 +24,7 @@ class NetgearDeviceEntity(CoordinatorEntity):
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator[Any],
+        coordinator: NetgearDataCoordinator[Any],
         router: NetgearRouter,
         device: dict,
     ) -> None:
@@ -90,12 +88,12 @@ class NetgearRouterEntity(Entity):
         )
 
 
-class NetgearRouterCoordinatorEntity(NetgearRouterEntity, CoordinatorEntity):
+class NetgearRouterCoordinatorEntity[T: NetgearDataCoordinator[Any]](
+    NetgearRouterEntity, CoordinatorEntity[T]
+):
     """Base class for a Netgear router entity."""
 
-    def __init__(
-        self, coordinator: DataUpdateCoordinator[Any], router: NetgearRouter
-    ) -> None:
+    def __init__(self, coordinator: T, router: NetgearRouter) -> None:
         """Initialize a Netgear device."""
         CoordinatorEntity.__init__(self, coordinator)
         NetgearRouterEntity.__init__(self, router)

homeassistant/components/netgear/sensor.py

@@ -26,9 +26,8 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.helpers.typing import StateType
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .coordinator import NetgearConfigEntry
+from .coordinator import NetgearConfigEntry, NetgearDataCoordinator
 from .entity import NetgearDeviceEntity, NetgearRouterCoordinatorEntity
 from .router import NetgearRouter
 
@@ -272,7 +271,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up Netgear sensors from a config entry."""
     router = entry.runtime_data.router
-    coordinator = entry.runtime_data.coordinator
+    coordinator_tracker = entry.runtime_data.coordinator_tracker
     coordinator_traffic = entry.runtime_data.coordinator_traffic
     coordinator_speed = entry.runtime_data.coordinator_speed
     coordinator_utilization = entry.runtime_data.coordinator_utilization
@@ -298,7 +297,7 @@ async def async_setup_entry(
     @callback
     def new_device_callback() -> None:
         """Add new devices if needed."""
-        if not coordinator.data:
+        if not coordinator_tracker.data:
             return
 
         new_entities: list[NetgearSensorEntity] = []
@@ -308,16 +307,16 @@ async def async_setup_entry(
                 continue
 
             new_entities.extend(
-                NetgearSensorEntity(coordinator, router, device, attribute)
+                NetgearSensorEntity(coordinator_tracker, router, device, attribute)
                 for attribute in sensors
             )
             tracked.add(mac)
 
         async_add_entities(new_entities)
 
-    entry.async_on_unload(coordinator.async_add_listener(new_device_callback))
+    entry.async_on_unload(coordinator_tracker.async_add_listener(new_device_callback))
 
-    coordinator.data = True
+    coordinator_tracker.data = True
     new_device_callback()
 
 
@@ -326,7 +325,7 @@ class NetgearSensorEntity(NetgearDeviceEntity, SensorEntity):
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator[Any],
+        coordinator: NetgearDataCoordinator[Any],
         router: NetgearRouter,
         device: dict,
         attribute: str,
@@ -365,7 +364,7 @@ class NetgearRouterSensorEntity(NetgearRouterCoordinatorEntity, RestoreSensor):
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator[dict[str, Any] | None],
+        coordinator: NetgearDataCoordinator[dict[str, Any] | None],
         router: NetgearRouter,
         entity_description: NetgearSensorEntityDescription,
     ) -> None:

homeassistant/components/netgear/switch.py

@@ -12,9 +12,8 @@ from homeassistant.components.switch import SwitchEntity, SwitchEntityDescriptio
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .coordinator import NetgearConfigEntry
+from .coordinator import NetgearConfigEntry, NetgearDataCoordinator
 from .entity import NetgearDeviceEntity, NetgearRouterEntity
 from .router import NetgearRouter
 
@@ -111,14 +110,14 @@ async def async_setup_entry(
     )
 
     # Entities per network device
-    coordinator = entry.runtime_data.coordinator
+    coordinator_tracker = entry.runtime_data.coordinator_tracker
     tracked = set()
 
     @callback
     def new_device_callback() -> None:
         """Add new devices if needed."""
         new_entities = []
-        if not coordinator.data:
+        if not coordinator_tracker.data:
             return
 
         for mac, device in router.devices.items():
@@ -127,7 +126,9 @@ async def async_setup_entry(
 
             new_entities.extend(
                 [
-                    NetgearAllowBlock(coordinator, router, device, entity_description)
+                    NetgearAllowBlock(
+                        coordinator_tracker, router, device, entity_description
+                    )
                     for entity_description in SWITCH_TYPES
                 ]
             )
@@ -135,9 +136,9 @@ async def async_setup_entry(
 
         async_add_entities(new_entities)
 
-    entry.async_on_unload(coordinator.async_add_listener(new_device_callback))
+    entry.async_on_unload(coordinator_tracker.async_add_listener(new_device_callback))
 
-    coordinator.data = True
+    coordinator_tracker.data = True
     new_device_callback()
 
 
@@ -148,7 +149,7 @@ class NetgearAllowBlock(NetgearDeviceEntity, SwitchEntity):
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator[bool],
+        coordinator: NetgearDataCoordinator[bool],
         router: NetgearRouter,
         device: dict,
         entity_description: SwitchEntityDescription,

homeassistant/components/netgear/update.py

@@ -12,9 +12,8 @@ from homeassistant.components.update import (
 )
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .coordinator import NetgearConfigEntry
+from .coordinator import NetgearConfigEntry, NetgearFirmwareCoordinator
 from .entity import NetgearRouterCoordinatorEntity
 from .router import NetgearRouter
 
@@ -34,7 +33,9 @@ async def async_setup_entry(
     async_add_entities(entities)
 
 
-class NetgearUpdateEntity(NetgearRouterCoordinatorEntity, UpdateEntity):
+class NetgearUpdateEntity(
+    NetgearRouterCoordinatorEntity[NetgearFirmwareCoordinator], UpdateEntity
+):
     """Update entity for a Netgear device."""
 
     _attr_device_class = UpdateDeviceClass.FIRMWARE
@@ -42,7 +43,7 @@ class NetgearUpdateEntity(NetgearRouterCoordinatorEntity, UpdateEntity):
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator[dict[str, Any] | None],
+        coordinator: NetgearFirmwareCoordinator,
         router: NetgearRouter,
     ) -> None:
         """Initialize a Netgear device."""