diff --git a/homeassistant/components/aussie_broadband/__init__.py b/homeassistant/components/aussie_broadband/__init__.py index 1fc7e47ebde..acae5b7fc44 100644 --- a/homeassistant/components/aussie_broadband/__init__.py +++ b/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() diff --git a/homeassistant/components/aussie_broadband/coordinator.py b/homeassistant/components/aussie_broadband/coordinator.py new file mode 100644 index 00000000000..7d53e664750 --- /dev/null +++ b/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 diff --git a/homeassistant/components/aussie_broadband/sensor.py b/homeassistant/components/aussie_broadband/sensor.py index 49796b3f6cd..b1f17c05679 100644 --- a/homeassistant/components/aussie_broadband/sensor.py +++ b/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