Implement coordinator for trafikverket_weather (#65233)

.coveragerc

@@ -1273,6 +1273,7 @@ omit =
     homeassistant/components/tradfri/switch.py
     homeassistant/components/trafikverket_train/sensor.py
     homeassistant/components/trafikverket_weatherstation/__init__.py
+    homeassistant/components/trafikverket_weatherstation/coordinator.py
     homeassistant/components/trafikverket_weatherstation/sensor.py
     homeassistant/components/transmission/sensor.py
     homeassistant/components/transmission/switch.py

homeassistant/components/trafikverket_weatherstation/__init__.py

@@ -1,22 +1,21 @@
 """The trafikverket_weatherstation component."""
 from __future__ import annotations
 
-import logging
-
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
-from .const import PLATFORMS
+from .const import DOMAIN, PLATFORMS
-
+from .coordinator import TVDataUpdateCoordinator
-_LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Trafikverket Weatherstation from a config entry."""
 
-    hass.config_entries.async_setup_platforms(entry, PLATFORMS)
+    coordinator = TVDataUpdateCoordinator(hass, entry)
+    await coordinator.async_config_entry_first_refresh()
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
 
-    _LOGGER.debug("Loaded entry for %s", entry.title)
+    hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
     return True
 
@@ -24,10 +23,4 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload Trafikverket Weatherstation config entry."""
 
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-
-    if unload_ok:
-        _LOGGER.debug("Unloaded entry for %s", entry.title)
-        return unload_ok
-
-    return False

homeassistant/components/trafikverket_weatherstation/const.py

@@ -5,8 +5,6 @@ DOMAIN = "trafikverket_weatherstation"
 CONF_STATION = "station"
 PLATFORMS = [Platform.SENSOR]
 ATTRIBUTION = "Data provided by Trafikverket"
-ATTR_MEASURE_TIME = "measure_time"
-ATTR_ACTIVE = "active"
 
 NONE_IS_ZERO_SENSORS = {
     "air_temp",

homeassistant/components/trafikverket_weatherstation/coordinator.py

@@ -0,0 +1,43 @@
+"""DataUpdateCoordinator for the Trafikverket Weather integration."""
+from __future__ import annotations
+
+from datetime import timedelta
+import logging
+
+from pytrafikverket.trafikverket_weather import TrafikverketWeather, WeatherStationInfo
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_API_KEY
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import CONF_STATION, DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+TIME_BETWEEN_UPDATES = timedelta(minutes=10)
+
+
+class TVDataUpdateCoordinator(DataUpdateCoordinator[WeatherStationInfo]):
+    """A Sensibo Data Update Coordinator."""
+
+    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
+        """Initialize the Sensibo coordinator."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            name=DOMAIN,
+            update_interval=TIME_BETWEEN_UPDATES,
+        )
+        self._weather_api = TrafikverketWeather(
+            async_get_clientsession(hass), entry.data[CONF_API_KEY]
+        )
+        self._station = entry.data[CONF_STATION]
+
+    async def _async_update_data(self) -> WeatherStationInfo:
+        """Fetch data from Trafikverket."""
+        try:
+            weatherdata = await self._weather_api.async_get_weather(self._station)
+        except ValueError as error:
+            raise UpdateFailed from error
+        return weatherdata

homeassistant/components/trafikverket_weatherstation/sensor.py

@@ -1,13 +1,8 @@
 """Weather information for air and road temperature (by Trafikverket)."""
 from __future__ import annotations
 
-import asyncio
 from dataclasses import dataclass
-from datetime import timedelta
+from datetime import datetime
-import logging
-
-import aiohttp
-from pytrafikverket.trafikverket_weather import TrafikverketWeather, WeatherStationInfo
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
@@ -17,7 +12,6 @@ from homeassistant.components.sensor import (
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
-    CONF_API_KEY,
     DEGREE,
     LENGTH_MILLIMETERS,
     PERCENTAGE,
@@ -25,26 +19,17 @@ from homeassistant.const import (
     TEMP_CELSIUS,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.device_registry import DeviceEntryType
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.util import Throttle
+from homeassistant.helpers.typing import StateType
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.util.dt import as_utc, get_time_zone
 
-from .const import (
+from .const import ATTRIBUTION, CONF_STATION, DOMAIN, NONE_IS_ZERO_SENSORS
-    ATTR_ACTIVE,
+from .coordinator import TVDataUpdateCoordinator
-    ATTR_MEASURE_TIME,
-    ATTRIBUTION,
-    CONF_STATION,
-    DOMAIN,
-    NONE_IS_ZERO_SENSORS,
-)
 
-_LOGGER = logging.getLogger(__name__)
+STOCKHOLM_TIMEZONE = get_time_zone("Europe/Stockholm")
-
-MIN_TIME_BETWEEN_UPDATES = timedelta(minutes=10)
-
-SCAN_INTERVAL = timedelta(seconds=300)
 
 
 @dataclass
@@ -143,6 +128,14 @@ SENSOR_TYPES: tuple[TrafikverketSensorEntityDescription, ...] = (
         icon="mdi:weather-pouring",
         entity_registry_enabled_default=False,
     ),
+    TrafikverketSensorEntityDescription(
+        key="measure_time",
+        api_key="measure_time",
+        name="Measure Time",
+        icon="mdi:clock",
+        entity_registry_enabled_default=False,
+        device_class=SensorDeviceClass.TIMESTAMP,
+    ),
 )
 
 
@@ -151,20 +144,25 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Trafikverket sensor entry."""
 
-    web_session = async_get_clientsession(hass)
+    coordinator: TVDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-    weather_api = TrafikverketWeather(web_session, entry.data[CONF_API_KEY])
 
-    entities = [
+    async_add_entities(
         TrafikverketWeatherStation(
-            weather_api, entry.entry_id, entry.data[CONF_STATION], description
+            coordinator, entry.entry_id, entry.data[CONF_STATION], description
         )
         for description in SENSOR_TYPES
-    ]
+    )
-
-    async_add_entities(entities, True)
 
 
-class TrafikverketWeatherStation(SensorEntity):
+def _to_datetime(measuretime: str) -> datetime:
+    """Return isoformatted utc time."""
+    time_obj = datetime.strptime(measuretime, "%Y-%m-%dT%H:%M:%S")
+    return as_utc(time_obj.replace(tzinfo=STOCKHOLM_TIMEZONE))
+
+
+class TrafikverketWeatherStation(
+    CoordinatorEntity[TVDataUpdateCoordinator], SensorEntity
+):
     """Representation of a Trafikverket sensor."""
 
     entity_description: TrafikverketSensorEntityDescription
@@ -172,17 +170,16 @@ class TrafikverketWeatherStation(SensorEntity):
 
     def __init__(
         self,
-        weather_api: TrafikverketWeather,
+        coordinator: TVDataUpdateCoordinator,
         entry_id: str,
         sensor_station: str,
         description: TrafikverketSensorEntityDescription,
     ) -> None:
         """Initialize the sensor."""
+        super().__init__(coordinator)
         self.entity_description = description
         self._attr_name = f"{sensor_station} {description.name}"
         self._attr_unique_id = f"{entry_id}_{description.key}"
-        self._station = sensor_station
-        self._weather_api = weather_api
         self._attr_device_info = DeviceInfo(
             entry_type=DeviceEntryType.SERVICE,
             identifiers={(DOMAIN, entry_id)},
@@ -191,26 +188,23 @@ class TrafikverketWeatherStation(SensorEntity):
             name=sensor_station,
             configuration_url="https://api.trafikinfo.trafikverket.se/",
         )
-        self._weather: WeatherStationInfo | None = None
 
-    @Throttle(MIN_TIME_BETWEEN_UPDATES)
+    @property
-    async def async_update(self) -> None:
+    def native_value(self) -> StateType | datetime:
-        """Get the latest data from Trafikverket and updates the states."""
+        """Return state of sensor."""
-        try:
+        if self.entity_description.api_key == "measure_time":
-            self._weather = await self._weather_api.async_get_weather(self._station)
+            return _to_datetime(self.coordinator.data.measure_time)
-        except (asyncio.TimeoutError, aiohttp.ClientError, ValueError) as error:
+
-            _LOGGER.error("Could not fetch weather data: %s", error)
+        state: StateType = getattr(
-            return
+            self.coordinator.data, self.entity_description.api_key
-        self._attr_native_value = getattr(
-            self._weather, self.entity_description.api_key
         )
-        if (
-            self._attr_native_value is None
-            and self.entity_description.key in NONE_IS_ZERO_SENSORS
-        ):
-            self._attr_native_value = 0
 
-        self._attr_extra_state_attributes = {
+        # For zero value state the api reports back None for certain sensors.
-            ATTR_ACTIVE: self._weather.active,
+        if state is None and self.entity_description.key in NONE_IS_ZERO_SENSORS:
-            ATTR_MEASURE_TIME: self._weather.measure_time,
+            return 0
-        }
+        return state
+
+    @property
+    def available(self) -> bool:
+        """Return if entity is available."""
+        return self.coordinator.data.active and super().available