Move aussie broadband coordinator to separate class (#127081)

homeassistant/components/aussie_broadband/__init__.py

@@ -2,24 +2,20 @@
 
 from __future__ import annotations
 
-from datetime import timedelta
-import logging
-
 from aiohttp import ClientError
 from aussiebb.asyncio import AussieBB
 from aussiebb.const import FETCH_TYPES
-from aussiebb.exceptions import AuthenticationException, UnrecognisedServiceType
+from aussiebb.exceptions import AuthenticationException
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .const import DEFAULT_UPDATE_INTERVAL, DOMAIN, SERVICE_ID
+from .const import DOMAIN
+from .coordinator import AussieBroadandDataUpdateCoordinator
 
-_LOGGER = logging.getLogger(__name__)
 PLATFORMS = [Platform.SENSOR]
 
 
@@ -43,24 +39,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     except ClientError as exc:
         raise ConfigEntryNotReady from exc
 
-    # Create an appropriate refresh function
-    def update_data_factory(service_id):
-        async def async_update_data():
-            try:
-                return await client.get_usage(service_id)
-            except UnrecognisedServiceType as err:
-                raise UpdateFailed(f"Service {service_id} was unrecognised") from err
-
-        return async_update_data
-
     # Initiate a Data Update Coordinator for each service
     for service in services:
-        service["coordinator"] = DataUpdateCoordinator(
-            hass,
-            _LOGGER,
-            name=service["service_id"],
-            update_interval=timedelta(minutes=DEFAULT_UPDATE_INTERVAL),
-            update_method=update_data_factory(service[SERVICE_ID]),
+        service["coordinator"] = AussieBroadandDataUpdateCoordinator(
+            hass, client, service["service_id"]
         )
         await service["coordinator"].async_config_entry_first_refresh()
 

homeassistant/components/aussie_broadband/coordinator.py

@@ -0,0 +1,39 @@
+"""Coordinator for the Aussie Broadband integration."""
+
+from __future__ import annotations
+
+from datetime import timedelta
+import logging
+from typing import Any
+
+from aussiebb.asyncio import AussieBB
+from aussiebb.exceptions import UnrecognisedServiceType
+
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import DEFAULT_UPDATE_INTERVAL
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class AussieBroadandDataUpdateCoordinator(DataUpdateCoordinator[dict[str, Any]]):
+    """Aussie Broadand data update coordinator."""
+
+    def __init__(self, hass: HomeAssistant, client: AussieBB, service_id: str) -> None:
+        """Initialize Atag coordinator."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            name=f"Aussie Broadband {service_id}",
+            update_interval=timedelta(minutes=DEFAULT_UPDATE_INTERVAL),
+        )
+        self._client = client
+        self._service_id = service_id
+
+    async def _async_update_data(self) -> dict[str, Any]:
+        """Update data via library."""
+        try:
+            return await self._client.get_usage(self._service_id)
+        except UnrecognisedServiceType as err:
+            raise UpdateFailed(f"Service {self._service_id} was unrecognised") from err

homeassistant/components/aussie_broadband/sensor.py

@@ -22,6 +22,7 @@ from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import DOMAIN, SERVICE_ID
+from .coordinator import AussieBroadandDataUpdateCoordinator
 
 
 @dataclass(frozen=True)
@@ -131,7 +132,9 @@ async def async_setup_entry(
     )
 
 
-class AussieBroadandSensorEntity(CoordinatorEntity, SensorEntity):
+class AussieBroadandSensorEntity(
+    CoordinatorEntity[AussieBroadandDataUpdateCoordinator], SensorEntity
+):
     """Base class for Aussie Broadband metric sensors."""
 
     _attr_has_entity_name = True