diff --git a/homeassistant/components/nsw_fuel_station/__init__.py b/homeassistant/components/nsw_fuel_station/__init__.py index 85e204b6f51..b1065d755f6 100644 --- a/homeassistant/components/nsw_fuel_station/__init__.py +++ b/homeassistant/components/nsw_fuel_station/__init__.py @@ -2,23 +2,16 @@ from __future__ import annotations -from dataclasses import dataclass -import datetime -import logging - -from nsw_fuel import FuelCheckClient, FuelCheckError, Station +from nsw_fuel import FuelCheckClient from homeassistant.core import HomeAssistant from homeassistant.helpers import config_validation as cv from homeassistant.helpers.typing import ConfigType -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed from .const import DATA_NSW_FUEL_STATION - -_LOGGER = logging.getLogger(__name__) +from .coordinator import NSWFuelStationCoordinator DOMAIN = "nsw_fuel_station" -SCAN_INTERVAL = datetime.timedelta(hours=1) CONFIG_SCHEMA = cv.platform_only_config_schema(DOMAIN) @@ -27,46 +20,9 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: """Set up the NSW Fuel Station platform.""" client = FuelCheckClient() - async def async_update_data(): - return await hass.async_add_executor_job(fetch_station_price_data, client) - - coordinator = DataUpdateCoordinator( - hass, - _LOGGER, - config_entry=None, - name="sensor", - update_interval=SCAN_INTERVAL, - update_method=async_update_data, - ) + coordinator = NSWFuelStationCoordinator(hass, client) hass.data[DATA_NSW_FUEL_STATION] = coordinator await coordinator.async_refresh() return True - - -@dataclass -class StationPriceData: - """Data structure for O(1) price and name lookups.""" - - stations: dict[int, Station] - prices: dict[tuple[int, str], float] - - -def fetch_station_price_data(client: FuelCheckClient) -> StationPriceData | None: - """Fetch fuel price and station data.""" - try: - raw_price_data = client.get_fuel_prices() - # Restructure prices and station details to be indexed by station code - # for O(1) lookup - return StationPriceData( - stations={s.code: s for s in raw_price_data.stations}, - prices={ - (p.station_code, p.fuel_type): p.price for p in raw_price_data.prices - }, - ) - - except FuelCheckError as exc: - raise UpdateFailed( - f"Failed to fetch NSW Fuel station price data: {exc}" - ) from exc diff --git a/homeassistant/components/nsw_fuel_station/coordinator.py b/homeassistant/components/nsw_fuel_station/coordinator.py new file mode 100644 index 00000000000..7856c8de5bd --- /dev/null +++ b/homeassistant/components/nsw_fuel_station/coordinator.py @@ -0,0 +1,65 @@ +"""Coordinator for the NSW Fuel Station integration.""" + +from __future__ import annotations + +from dataclasses import dataclass +import datetime +import logging + +from nsw_fuel import FuelCheckClient, FuelCheckError, Station + +from homeassistant.core import HomeAssistant +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed + +_LOGGER = logging.getLogger(__name__) + +SCAN_INTERVAL = datetime.timedelta(hours=1) + + +@dataclass +class StationPriceData: + """Data structure for O(1) price and name lookups.""" + + stations: dict[int, Station] + prices: dict[tuple[int, str], float] + + +class NSWFuelStationCoordinator(DataUpdateCoordinator[StationPriceData | None]): + """Class to manage fetching NSW fuel station data.""" + + config_entry: None + + def __init__(self, hass: HomeAssistant, client: FuelCheckClient) -> None: + """Initialize the coordinator.""" + super().__init__( + hass, + _LOGGER, + config_entry=None, + name="sensor", + update_interval=SCAN_INTERVAL, + ) + self.client = client + + async def _async_update_data(self) -> StationPriceData | None: + """Fetch data from API.""" + return await self.hass.async_add_executor_job( + _fetch_station_price_data, self.client + ) + + +def _fetch_station_price_data(client: FuelCheckClient) -> StationPriceData | None: + """Fetch fuel price and station data.""" + try: + raw_price_data = client.get_fuel_prices() + # Restructure prices and station details to be indexed by station code + # for O(1) lookup + return StationPriceData( + stations={s.code: s for s in raw_price_data.stations}, + prices={ + (p.station_code, p.fuel_type): p.price for p in raw_price_data.prices + }, + ) + except FuelCheckError as exc: + raise UpdateFailed( + f"Failed to fetch NSW Fuel station price data: {exc}" + ) from exc diff --git a/homeassistant/components/nsw_fuel_station/sensor.py b/homeassistant/components/nsw_fuel_station/sensor.py index 7ae9b3a4d9f..81c5d4d070f 100644 --- a/homeassistant/components/nsw_fuel_station/sensor.py +++ b/homeassistant/components/nsw_fuel_station/sensor.py @@ -15,12 +15,10 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers import config_validation as cv from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType -from homeassistant.helpers.update_coordinator import ( - CoordinatorEntity, - DataUpdateCoordinator, -) +from homeassistant.helpers.update_coordinator import CoordinatorEntity -from . import DATA_NSW_FUEL_STATION, StationPriceData +from .const import DATA_NSW_FUEL_STATION +from .coordinator import NSWFuelStationCoordinator _LOGGER = logging.getLogger(__name__) @@ -65,7 +63,7 @@ def setup_platform( station_id = config[CONF_STATION_ID] fuel_types = config[CONF_FUEL_TYPES] - coordinator = hass.data[DATA_NSW_FUEL_STATION] + coordinator: NSWFuelStationCoordinator = hass.data[DATA_NSW_FUEL_STATION] if coordinator.data is None: _LOGGER.error("Initial fuel station price data not available") @@ -86,16 +84,14 @@ def setup_platform( add_entities(entities) -class StationPriceSensor( - CoordinatorEntity[DataUpdateCoordinator[StationPriceData]], SensorEntity -): +class StationPriceSensor(CoordinatorEntity[NSWFuelStationCoordinator], SensorEntity): """Implementation of a sensor that reports the fuel price for a station.""" _attr_attribution = "Data provided by NSW Government FuelCheck" def __init__( self, - coordinator: DataUpdateCoordinator[StationPriceData], + coordinator: NSWFuelStationCoordinator, station_id: int, fuel_type: str, ) -> None: