Refactor eafm to avoid creating entities in the coordinator update (#111601)

homeassistant/components/eafm/__init__.py

@@ -1,16 +1,58 @@
 """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
 
 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:
     """Set up flood monitoring sensors for this config entry."""
-    hass.data.setdefault(DOMAIN, {})
+    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,
+        name="sensor",
+        update_method=_async_update_data,
+        update_interval=timedelta(seconds=15 * 60),
+    )
+    await coordinator.async_config_entry_first_refresh()
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     return True
 

homeassistant/components/eafm/sensor.py

@@ -1,15 +1,11 @@
 """Support for gauges from flood monitoring API."""
-import asyncio
-from datetime import timedelta
-import logging
 
-from aioeafm import get_station
+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
+from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import (
@@ -19,72 +15,49 @@ from homeassistant.helpers.update_coordinator import (
 
 from .const import DOMAIN
 
-_LOGGER = logging.getLogger(__name__)
-
 UNIT_MAPPING = {
     "http://qudt.org/1.1/vocab/unit#Meter": UnitOfLength.METERS,
 }
 
 
-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,
     config_entry: ConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up UK Flood Monitoring Sensors."""
-    station_key = config_entry.data["station"]
+    coordinator: DataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
-    session = async_get_clientsession(hass=hass)
+    created_entities: set[str] = set()
-
-    measurements = set()
-
-    async def async_update_data():
-        # DataUpdateCoordinator will handle aiohttp ClientErrors and timeouts
-        async with asyncio.timeout(30):
-            data = await get_station(session, station_key)
-
-        measures = get_measures(data)
-        entities = []
 
+    @callback
+    def _async_create_new_entities():
+        """Create new entities."""
+        if not coordinator.last_update_success:
+            return
+        measures: dict[str, dict[str, Any]] = coordinator.data["measures"]
+        entities: list[Measurement] = []
         # Look to see if payload contains new measures
-        for measure in measures:
+        for key, data in measures.items():
-            if measure["@id"] in measurements:
+            if key in created_entities:
                 continue
 
-            if "latestReading" not in measure:
+            if "latestReading" not in data:
                 # Don't create a sensor entity for a gauge that isn't available
                 continue
 
-            entities.append(Measurement(hass.data[DOMAIN][station_key], measure["@id"]))
+            entities.append(Measurement(coordinator, key))
-            measurements.add(measure["@id"])
+            created_entities.add(key)
 
         async_add_entities(entities)
 
-        # Turn data.measures into a dict rather than a list so easier for entities to
+    _async_create_new_entities()
-        # find themselves.
-        data["measures"] = {measure["@id"]: measure for measure in measures}
 
-        return data
+    # Subscribe to the coordinator to create new entities
-
+    # when the coordinator updates
-    hass.data[DOMAIN][station_key] = coordinator = DataUpdateCoordinator(
+    config_entry.async_on_unload(
-        hass,
+        coordinator.async_add_listener(_async_create_new_entities)
-        _LOGGER,
-        name="sensor",
-        update_method=async_update_data,
-        update_interval=timedelta(seconds=15 * 60),
     )
 
-    # Fetch initial data so we have data when entities subscribe
-    await coordinator.async_refresh()
-
 
 class Measurement(CoordinatorEntity, SensorEntity):
     """A gauge at a flood monitoring station."""

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.sensor.get_station") as patched:
+    with patch("homeassistant.components.eafm.get_station") as patched:
         yield patched

tests/components/eafm/test_sensor.py

@@ -339,6 +339,19 @@ async def test_ignore_no_latest_reading(hass: HomeAssistant, mock_get_station) -
     assert state is None
 
 
+async def test_no_measures(hass: HomeAssistant, mock_get_station) -> None:
+    """Test no measures in the data."""
+    await async_setup_test_fixture(
+        hass,
+        mock_get_station,
+        {
+            "label": "My station",
+        },
+    )
+
+    assert hass.states.async_entity_ids_count() == 0
+
+
 async def test_mark_existing_as_unavailable_if_no_latest(
     hass: HomeAssistant, mock_get_station
 ) -> None: