Migrate nsw_fuel to DataCoordinator

home-assistant · buxtronix · May 22, 2025 · May 26, 2025 · Jun 1, 2025 · Jun 2, 2025
commit 5540a25c43488c207b2905b0b09b9227dbd122f3
diff --git a/homeassistant/components/nsw_fuel_station/__init__.py b/homeassistant/components/nsw_fuel_station/__init__.py
@@ -4,23 +4,22 @@
 
 from dataclasses import dataclass
 import datetime
-import logging
 from typing import TYPE_CHECKING
 
-from nsw_fuel import FuelCheckClient, FuelCheckError, Station
+from nsw_fuel import Station
 
 from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.helpers import config_validation as cv
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import DOMAIN
+from .coordinator import NswFuelStationDataUpdateCoordinator
 
 if TYPE_CHECKING:
     from homeassistant.core import HomeAssistant
     from homeassistant.helpers.typing import ConfigType
 
-from .const import DATA_NSW_FUEL_STATION
 
-DOMAIN = "nsw_fuel_station"
 SCAN_INTERVAL = datetime.timedelta(hours=1)
 
 CONFIG_SCHEMA = cv.platform_only_config_schema(DOMAIN)
@@ -29,33 +28,32 @@
     Platform.SENSOR,
 ]
 
-_LOGGER = logging.getLogger(__name__)
 
+class FuelCheckData:
+    """Holds a global coordinator."""
 
-async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
-    """Set up the NSW Fuel Station platform."""
-    client = FuelCheckClient()
-
-    async def async_update_data() -> StationPriceData:
-        return await hass.async_add_executor_job(fetch_station_price_data, hass, client)
+    def __init__(self, hass: HomeAssistant) -> None:
+        """Initialise the data."""
+        self.hass: HomeAssistant = hass
+        self.coordinator: NswFuelStationDataUpdateCoordinator = (
+            NswFuelStationDataUpdateCoordinator(
+                hass,
+            )
+        )
 
-    coordinator = DataUpdateCoordinator(
-        hass,
-        _LOGGER,
-        config_entry=None,
-        name="sensor",
-        update_interval=SCAN_INTERVAL,
-        update_method=async_update_data,
-    )
-    hass.data[DATA_NSW_FUEL_STATION] = coordinator
 
-    await coordinator.async_refresh()
+async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+    """Set up the NSW Fuel Station platform."""
+    fuel_data = hass.data.setdefault(DOMAIN, {})
+    if "coordinator" not in fuel_data:
+        fuel_data["coordinator"] = FuelCheckData(hass).coordinator
+    await fuel_data["coordinator"].async_config_entry_first_refresh()
 
     if "sensor" not in config:
         return True
 
     for platform_config in config["sensor"]:
-        if platform_config["platform"] == "nsw_fuel_station":
+        if platform_config["platform"] == DOMAIN:
             hass.async_create_task(
                 hass.config_entries.flow.async_init(
                     DOMAIN,
@@ -72,6 +70,10 @@ async def async_setup_entry(
     entry: ConfigEntry,
 ) -> bool:
     """Set up this integration using UI."""
+    fuel_data = hass.data.setdefault(DOMAIN, {})
+    if "coordinator" not in fuel_data:
+        fuel_data["coordinator"] = FuelCheckData(hass).coordinator
+    await fuel_data["coordinator"].async_config_entry_first_refresh()
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True
@@ -92,28 +94,3 @@ class StationPriceData:
     stations: dict[int, Station]
     prices: dict[tuple[int, str], float]
     fuel_types: dict[str, str]
-
-
-def fetch_station_price_data(
-    hass: HomeAssistant, client: FuelCheckClient
-) -> StationPriceData:
-    """Fetch fuel price and station data."""
-    try:
-        raw_price_data = client.get_fuel_prices()
-        reference_data = client.get_reference_data()
-    except FuelCheckError as exc:
-        raise UpdateFailed(
-            f"Failed to fetch NSW Fuel station price data: {exc}"
-        ) from exc
-
-    # Restructure prices and station details to be indexed by station code
-    # for O(1) lookup
-    price_data = 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},
-        fuel_types={},
-    )
-
-    price_data.fuel_types = {f.code: f.name for f in reference_data.fuel_types}
-
-    return price_data
diff --git a/homeassistant/components/nsw_fuel_station/config_flow.py b/homeassistant/components/nsw_fuel_station/config_flow.py
@@ -2,10 +2,9 @@
 
 from __future__ import annotations
 
-import logging
 from typing import Any
 
-from nsw_fuel import FuelCheckClient, Station
+from nsw_fuel import Station
 import voluptuous as vol
 
 from homeassistant.config_entries import (
@@ -19,22 +18,20 @@
     SelectSelectorConfig,
     SelectSelectorMode,
 )
-from homeassistant.helpers.update_coordinator import UpdateFailed
 
-from . import StationPriceData, fetch_station_price_data
+from . import StationPriceData
 from .const import (
     CONF_FUEL_TYPES,
     CONF_STATION_ID,
-    DATA_NSW_FUEL_STATION,
+    DOMAIN,
     INPUT_FUEL_TYPES,
     INPUT_SEARCH_TERM,
     INPUT_STATION_ID,
 )
+from .coordinator import NswFuelStationDataUpdateCoordinator
 
-_LOGGER = logging.getLogger(__name__)
 
-
-class FuelCheckConfigFlow(ConfigFlow, domain=DATA_NSW_FUEL_STATION):
+class FuelCheckConfigFlow(ConfigFlow, domain=DOMAIN):
     """Handle a config flow for Fuel Check integration."""
 
     VERSION = 1
@@ -46,18 +43,21 @@ class FuelCheckConfigFlow(ConfigFlow, domain=DATA_NSW_FUEL_STATION):
     selected_station: Station
     data: StationPriceData
 
+    async def _setup_coordinator(self):
+        coordinator = self.hass.data.get(DOMAIN, {}).get("coordinator")
+        if coordinator is None:
+            coordinator = NswFuelStationDataUpdateCoordinator(self.hass)
+            await coordinator.async_config_entry_first_refresh()
+        self.data = coordinator.data
+
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Display the initial UI form."""
         errors: dict[str, str] = {}
 
-        try:
-            self.data = await self.hass.async_add_executor_job(
-                fetch_station_price_data, self.hass, FuelCheckClient()
-            )
-        except UpdateFailed as e:
-            _LOGGER.error("Failed to fetch fuel price data: %s", e)
+        await self._setup_coordinator()
+        if self.data is None:
             return self.async_abort(reason="fetch_failed")
 
         if user_input is not None:
@@ -185,6 +185,11 @@ async def async_step_import(
             CONF_FUEL_TYPES: user_input[INPUT_FUEL_TYPES],
         }
         self._async_abort_entries_match({CONF_STATION_ID: station_id})
+
+        await self._setup_coordinator()
+        if self.data is None:
+            return self.async_abort(reason="fetch_failed")
+
         station = self.data.stations.get(station_id)
         name = "Unknown"
         if station is not None:

diff --git a/homeassistant/components/nsw_fuel_station/const.py b/homeassistant/components/nsw_fuel_station/const.py
@@ -1,6 +1,6 @@
 """Constants for the NSW Fuel Station integration."""
 
-DATA_NSW_FUEL_STATION = "nsw_fuel_station"
+DOMAIN = "nsw_fuel_station"
 
 INPUT_FUEL_TYPES = "fuel_types"
 INPUT_STATION_ID = "station_id"

diff --git a/homeassistant/components/nsw_fuel_station/coordinator.py b/homeassistant/components/nsw_fuel_station/coordinator.py
@@ -0,0 +1,75 @@
+"""Data coordinator for nsw_fuel."""
+
+from dataclasses import dataclass
+from datetime import timedelta
+import logging
+
+from nsw_fuel import FuelCheckClient, Station
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
+
+from .const import DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+
+type NswFuelStationConfigEntry = ConfigEntry[NswFuelStationDataUpdateCoordinator]
+
+
+@dataclass
+class NswFuelStationCoordinatorData:
+    """Data class for storing coordinator data."""
+
+    stations: dict[int, Station]
+    prices: dict[tuple[int, str], float]
+    fuel_types: dict[str, str]
+
+
+class NswFuelStationDataUpdateCoordinator(
+    DataUpdateCoordinator[NswFuelStationCoordinatorData]
+):
+    """Coordinator for Nsw Fuel data."""
+
+    client: FuelCheckClient
+    config_entry: NswFuelStationConfigEntry
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+    ) -> None:
+        """Initialise the data coordinator."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            name=DOMAIN,
+            update_interval=timedelta(hours=1),
+        )
+        self.client = FuelCheckClient()
+
+    async def _async_update_data(self) -> NswFuelStationCoordinatorData:
+        """Fetch updated data."""
+        price_data = NswFuelStationCoordinatorData(
+            stations={},
+            prices={},
+            fuel_types={},
+        )
+        raw_price_data = await self.hass.async_add_executor_job(
+            self.client.get_fuel_prices
+        )
+        if self.data is None or len(self.data.fuel_types) < 1:
+            reference_data = await self.hass.async_add_executor_job(
+                self.client.get_reference_data
+            )
+            price_data.fuel_types = {f.code: f.name for f in reference_data.fuel_types}
+        else:
+            price_data.fuel_types = self.data.fuel_types
+
+        # Restructure prices and station details to be indexed by station code
+        # for O(1) lookup
+        price_data.stations = {s.code: s for s in raw_price_data.stations}
+        price_data.prices = {
+            (p.sta
C799
tion_code, p.fuel_type): p.price for p in raw_price_data.prices
+        }
+
+        return price_data
diff --git a/homeassistant/components/nsw_fuel_station/sensor.py b/homeassistant/components/nsw_fuel_station/sensor.py
@@ -14,14 +14,15 @@
     DataUpdateCoordinator,
 )
 
-from . import DATA_NSW_FUEL_STATION, StationPriceData
+from . import StationPriceData
 from .const import (
     ATTR_FUEL_TYPE,
     ATTR_STATION_ADDRESS,
     ATTR_STATION_ID,
     ATTR_STATION_NAME,
     CONF_FUEL_TYPES,
     CONF_STATION_ID,
+    DOMAIN,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -37,7 +38,7 @@ async def async_setup_entry(
     station_id = entry.data[CONF_STATION_ID]
     fuel_types = entry.data[CONF_FUEL_TYPES]
 
-    coordinator = hass.data[DATA_NSW_FUEL_STATION]
+    coordinator = hass.data[DOMAIN]["coordinator"]
 
     if coordinator.data is None:
         _LOGGER.error("Initial fuel station price data not available")