From 7dc2b92452970da84b66c38261642cc72bc9aeb8 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Mon, 27 Jan 2025 15:05:20 +0100 Subject: [PATCH] Use typed coordinator and runtime_data in eafm (#136629) * Move coordinator and use runtime_data in eafm * Add type hints --- homeassistant/components/eafm/__init__.py | 52 ++---------------- homeassistant/components/eafm/coordinator.py | 57 ++++++++++++++++++++ homeassistant/components/eafm/sensor.py | 13 ++--- tests/components/eafm/conftest.py | 2 +- 4 files changed, 68 insertions(+), 56 deletions(-) create mode 100644 homeassistant/components/eafm/coordinator.py diff --git a/homeassistant/components/eafm/__init__.py b/homeassistant/components/eafm/__init__.py index dc618a983f3..e2af2bae9f5 100644 --- a/homeassistant/components/eafm/__init__.py +++ b/homeassistant/components/eafm/__init__.py @@ -1,64 +1,22 @@ """UK Environment Agency Flood Monitoring Integration.""" -import asyncio -from datetime import timedelta -import logging -from typing import Any - -from aioeafm import get_station - -from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant -from homeassistant.helpers.aiohttp_client import async_get_clientsession -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from .const import DOMAIN +from .coordinator import EafmConfigEntry, EafmCoordinator PLATFORMS = [Platform.SENSOR] -_LOGGER = logging.getLogger(__name__) - -def get_measures(station_data): - """Force measure key to always be a list.""" - if "measures" not in station_data: - return [] - if isinstance(station_data["measures"], dict): - return [station_data["measures"]] - return station_data["measures"] - - -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: EafmConfigEntry) -> bool: """Set up flood monitoring sensors for this config entry.""" - station_key = entry.data["station"] - session = async_get_clientsession(hass=hass) - - async def _async_update_data() -> dict[str, dict[str, Any]]: - # DataUpdateCoordinator will handle aiohttp ClientErrors and timeouts - async with asyncio.timeout(30): - data = await get_station(session, station_key) - - measures = get_measures(data) - # Turn data.measures into a dict rather than a list so easier for entities to - # find themselves. - data["measures"] = {measure["@id"]: measure for measure in measures} - return data - - coordinator = DataUpdateCoordinator[dict[str, dict[str, Any]]]( - hass, - _LOGGER, - config_entry=entry, - name="sensor", - update_method=_async_update_data, - update_interval=timedelta(seconds=15 * 60), - ) + coordinator = EafmCoordinator(hass, entry=entry) await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator + entry.runtime_data = coordinator await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: EafmConfigEntry) -> bool: """Unload flood monitoring sensors.""" return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/eafm/coordinator.py b/homeassistant/components/eafm/coordinator.py new file mode 100644 index 00000000000..375368210a5 --- /dev/null +++ b/homeassistant/components/eafm/coordinator.py @@ -0,0 +1,57 @@ +"""UK Environment Agency Flood Monitoring Integration.""" + +import asyncio +from datetime import timedelta +import logging +from typing import Any + +from aioeafm import get_station + +from homeassistant.config_entries import ConfigEntry +from homeassistant.const import Platform +from homeassistant.core import HomeAssistant +from homeassistant.helpers.aiohttp_client import async_get_clientsession +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator + +PLATFORMS = [Platform.SENSOR] + +_LOGGER = logging.getLogger(__name__) + +type EafmConfigEntry = ConfigEntry[EafmCoordinator] + + +def _get_measures(station_data: dict[str, Any]) -> list[dict[str, Any]]: + """Force measure key to always be a list.""" + if "measures" not in station_data: + return [] + if isinstance(station_data["measures"], dict): + return [station_data["measures"]] + return station_data["measures"] + + +class EafmCoordinator(DataUpdateCoordinator[dict[str, dict[str, Any]]]): + """Class to manage fetching UK Flood Monitoring data.""" + + def __init__(self, hass: HomeAssistant, entry: EafmConfigEntry) -> None: + """Initialize.""" + self._station_key = entry.data["station"] + self._session = async_get_clientsession(hass=hass) + super().__init__( + hass, + _LOGGER, + config_entry=entry, + name="sensor", + update_interval=timedelta(seconds=15 * 60), + ) + + async def _async_update_data(self) -> dict[str, dict[str, Any]]: + """Fetch the latest data from the source.""" + # DataUpdateCoordinator will handle aiohttp ClientErrors and timeouts + async with asyncio.timeout(30): + data = await get_station(self._session, self._station_key) + + measures = _get_measures(data) + # Turn data.measures into a dict rather than a list so easier for entities to + # find themselves. + data["measures"] = {measure["@id"]: measure for measure in measures} + return data diff --git a/homeassistant/components/eafm/sensor.py b/homeassistant/components/eafm/sensor.py index 297f4d6d2c8..d9b18cbc663 100644 --- a/homeassistant/components/eafm/sensor.py +++ b/homeassistant/components/eafm/sensor.py @@ -3,17 +3,14 @@ from typing import Any from homeassistant.components.sensor import SensorEntity, SensorStateClass -from homeassistant.config_entries import ConfigEntry from homeassistant.const import UnitOfLength from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback -from homeassistant.helpers.update_coordinator import ( - CoordinatorEntity, - DataUpdateCoordinator, -) +from homeassistant.helpers.update_coordinator import CoordinatorEntity from .const import DOMAIN +from .coordinator import EafmConfigEntry, EafmCoordinator UNIT_MAPPING = { "http://qudt.org/1.1/vocab/unit#Meter": UnitOfLength.METERS, @@ -22,11 +19,11 @@ UNIT_MAPPING = { async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: EafmConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up UK Flood Monitoring Sensors.""" - coordinator: DataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id] + coordinator = config_entry.runtime_data created_entities: set[str] = set() @callback @@ -70,7 +67,7 @@ class Measurement(CoordinatorEntity, SensorEntity): _attr_has_entity_name = True _attr_name = None - def __init__(self, coordinator, key): + def __init__(self, coordinator: EafmCoordinator, key: str) -> None: """Initialise the gauge with a data instance and station.""" super().__init__(coordinator) self.key = key diff --git a/tests/components/eafm/conftest.py b/tests/components/eafm/conftest.py index 3b060563a30..5dbdc98ad29 100644 --- a/tests/components/eafm/conftest.py +++ b/tests/components/eafm/conftest.py @@ -15,5 +15,5 @@ def mock_get_stations(): @pytest.fixture def mock_get_station(): """Mock aioeafm.get_station.""" - with patch("homeassistant.components.eafm.get_station") as patched: + with patch("homeassistant.components.eafm.coordinator.get_station") as patched: yield patched