Add coordinator to SMHI (#139052)

* Add coordinator to SMHI

* Remove not needed logging

* docstrings

homeassistant/components/smhi/__init__.py

@@ -1,6 +1,5 @@
 """Support for the Swedish weather institute weather service."""
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_LATITUDE,
     CONF_LOCATION,
@@ -10,10 +9,12 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant
 
+from .coordinator import SMHIConfigEntry, SMHIDataUpdateCoordinator
+
 PLATFORMS = [Platform.WEATHER]
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: SMHIConfigEntry) -> bool:
     """Set up SMHI forecast as config entry."""
 
     # Setting unique id where missing
@@ -21,16 +22,20 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         unique_id = f"{entry.data[CONF_LOCATION][CONF_LATITUDE]}-{entry.data[CONF_LOCATION][CONF_LONGITUDE]}"
         hass.config_entries.async_update_entry(entry, unique_id=unique_id)
 
+    coordinator = SMHIDataUpdateCoordinator(hass, entry)
+    await coordinator.async_config_entry_first_refresh()
+    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: SMHIConfigEntry) -> bool:
     """Unload a config entry."""
     return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 
-async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_migrate_entry(hass: HomeAssistant, entry: SMHIConfigEntry) -> bool:
     """Migrate old entry."""
 
     if entry.version > 3:

homeassistant/components/smhi/const.py

@@ -1,5 +1,7 @@
 """Constants in smhi component."""
 
+from datetime import timedelta
+import logging
 from typing import Final
 
 from homeassistant.components.weather import DOMAIN as WEATHER_DOMAIN
@@ -12,3 +14,8 @@ HOME_LOCATION_NAME = "Home"
 DEFAULT_NAME = "Weather"
 
 ENTITY_ID_SENSOR_FORMAT = WEATHER_DOMAIN + ".smhi_{}"
+
+LOGGER = logging.getLogger(__package__)
+
+DEFAULT_SCAN_INTERVAL = timedelta(minutes=31)
+TIMEOUT = 10

homeassistant/components/smhi/coordinator.py

@@ -0,0 +1,63 @@
+"""DataUpdateCoordinator for the SMHI integration."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+
+from pysmhi import SMHIForecast, SmhiForecastException, SMHIPointForecast
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_LATITUDE, CONF_LOCATION, CONF_LONGITUDE
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import aiohttp_client
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import DEFAULT_SCAN_INTERVAL, DOMAIN, LOGGER, TIMEOUT
+
+type SMHIConfigEntry = ConfigEntry[SMHIDataUpdateCoordinator]
+
+
+@dataclass
+class SMHIForecastData:
+    """Dataclass for SMHI data."""
+
+    daily: list[SMHIForecast]
+    hourly: list[SMHIForecast]
+
+
+class SMHIDataUpdateCoordinator(DataUpdateCoordinator[SMHIForecastData]):
+    """A SMHI Data Update Coordinator."""
+
+    config_entry: SMHIConfigEntry
+
+    def __init__(self, hass: HomeAssistant, config_entry: SMHIConfigEntry) -> None:
+        """Initialize the SMHI coordinator."""
+        super().__init__(
+            hass,
+            LOGGER,
+            config_entry=config_entry,
+            name=DOMAIN,
+            update_interval=DEFAULT_SCAN_INTERVAL,
+        )
+        self._smhi_api = SMHIPointForecast(
+            config_entry.data[CONF_LOCATION][CONF_LONGITUDE],
+            config_entry.data[CONF_LOCATION][CONF_LATITUDE],
+            session=aiohttp_client.async_get_clientsession(hass),
+        )
+
+    async def _async_update_data(self) -> SMHIForecastData:
+        """Fetch data from SMHI."""
+        try:
+            async with asyncio.timeout(TIMEOUT):
+                _forecast_daily = await self._smhi_api.async_get_daily_forecast()
+                _forecast_hourly = await self._smhi_api.async_get_hourly_forecast()
+        except SmhiForecastException as ex:
+            raise UpdateFailed(
+                "Failed to retrieve the forecast from the SMHI API"
+            ) from ex
+
+        return SMHIForecastData(
+            daily=_forecast_daily,
+            hourly=_forecast_hourly,
+        )

homeassistant/components/smhi/entity.py

@@ -2,16 +2,16 @@
 
 from __future__ import annotations
 
-import aiohttp
+from abc import abstractmethod
-from pysmhi import SMHIPointForecast
 
 from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
-from homeassistant.helpers.entity import Entity
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import DOMAIN
+from .coordinator import SMHIDataUpdateCoordinator
 
 
-class SmhiWeatherBaseEntity(Entity):
+class SmhiWeatherBaseEntity(CoordinatorEntity[SMHIDataUpdateCoordinator]):
     """Representation of a base weather entity."""
 
     _attr_attribution = "Swedish weather institute (SMHI)"
@@ -22,11 +22,11 @@ class SmhiWeatherBaseEntity(Entity):
         self,
         latitude: str,
         longitude: str,
-        session: aiohttp.ClientSession,
+        coordinator: SMHIDataUpdateCoordinator,
     ) -> None:
         """Initialize the SMHI base weather entity."""
+        super().__init__(coordinator)
         self._attr_unique_id = f"{latitude}, {longitude}"
-        self._smhi_api = SMHIPointForecast(longitude, latitude, session=session)
         self._attr_device_info = DeviceInfo(
             entry_type=DeviceEntryType.SERVICE,
             identifiers={(DOMAIN, f"{latitude}, {longitude}")},
@@ -34,3 +34,8 @@ class SmhiWeatherBaseEntity(Entity):
             model="v2",
             configuration_url="http://opendata.smhi.se/apidocs/metfcst/parameters.html",
         )
+        self.update_entity_data()
+
+    @abstractmethod
+    def update_entity_data(self) -> None:
+        """Refresh the entity data."""

homeassistant/components/smhi/weather.py

@@ -2,14 +2,11 @@
 
 from __future__ import annotations
 
-import asyncio
 from collections.abc import Mapping
-from datetime import datetime, timedelta
+from datetime import timedelta
-import logging
 from typing import Any, Final
 
-import aiohttp
+from pysmhi import SMHIForecast
-from pysmhi import SMHIForecast, SmhiForecastException
 
 from homeassistant.components.weather import (
     ATTR_CONDITION_CLEAR_NIGHT,
@@ -39,10 +36,9 @@ from homeassistant.components.weather import (
     ATTR_FORECAST_TIME,
     ATTR_FORECAST_WIND_BEARING,
     Forecast,
-    WeatherEntity,
+    SingleCoordinatorWeatherEntity,
     WeatherEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_LATITUDE,
     CONF_LOCATION,
@@ -53,17 +49,14 @@ from homeassistant.const import (
     UnitOfSpeed,
     UnitOfTemperature,
 )
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import aiohttp_client, sun
+from homeassistant.helpers import sun
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
-from homeassistant.helpers.event import async_call_later
-from homeassistant.util import Throttle
 
 from .const import ATTR_SMHI_THUNDER_PROBABILITY, ENTITY_ID_SENSOR_FORMAT
+from .coordinator import SMHIConfigEntry
 from .entity import SmhiWeatherBaseEntity
 
-_LOGGER = logging.getLogger(__name__)
-
 # Used to map condition from API results
 CONDITION_CLASSES: Final[dict[str, list[int]]] = {
     ATTR_CONDITION_CLOUDY: [5, 6],
@@ -96,25 +89,25 @@ MIN_TIME_BETWEEN_UPDATES = timedelta(minutes=31)
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: SMHIConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Add a weather entity from map location."""
     location = config_entry.data
 
-    session = aiohttp_client.async_get_clientsession(hass)
+    coordinator = config_entry.runtime_data
 
     entity = SmhiWeather(
         location[CONF_LOCATION][CONF_LATITUDE],
         location[CONF_LOCATION][CONF_LONGITUDE],
-        session=session,
+        coordinator=coordinator,
     )
     entity.entity_id = ENTITY_ID_SENSOR_FORMAT.format(config_entry.title)
 
-    async_add_entities([entity], True)
+    async_add_entities([entity])
 
 
-class SmhiWeather(SmhiWeatherBaseEntity, WeatherEntity):
+class SmhiWeather(SmhiWeatherBaseEntity, SingleCoordinatorWeatherEntity):
     """Representation of a weather entity."""
 
     _attr_native_temperature_unit = UnitOfTemperature.CELSIUS
@@ -126,61 +119,37 @@ class SmhiWeather(SmhiWeatherBaseEntity, WeatherEntity):
         WeatherEntityFeature.FORECAST_DAILY | WeatherEntityFeature.FORECAST_HOURLY
     )
 
-    def __init__(
+    def update_entity_data(self) -> None:
-        self,
+        """Refresh the entity data."""
-        latitude: str,
+        if daily_data := self.coordinator.data.daily:
-        longitude: str,
+            self._attr_native_temperature = daily_data[0]["temperature"]
-        session: aiohttp.ClientSession,
+            self._attr_humidity = daily_data[0]["humidity"]
-    ) -> None:
+            self._attr_native_wind_speed = daily_data[0]["wind_speed"]
-        """Initialize the SMHI weather entity."""
+            self._attr_wind_bearing = daily_data[0]["wind_direction"]
-        super().__init__(latitude, longitude, session)
+            self._attr_native_visibility = daily_data[0]["visibility"]
-        self._forecast_daily: list[SMHIForecast] | None = None
+            self._attr_native_pressure = daily_data[0]["pressure"]
-        self._forecast_hourly: list[SMHIForecast] | None = None
+            self._attr_native_wind_gust_speed = daily_data[0]["wind_gust"]
-        self._fail_count = 0
+            self._attr_cloud_coverage = daily_data[0]["total_cloud"]
+            self._attr_condition = CONDITION_MAP.get(daily_data[0]["symbol"])
+            if self._attr_condition == ATTR_CONDITION_SUNNY and not sun.is_up(
+                self.coordinator.hass
+            ):
+                self._attr_condition = ATTR_CONDITION_CLEAR_NIGHT
 
     @property
     def extra_state_attributes(self) -> Mapping[str, Any] | None:
         """Return additional attributes."""
-        if self._forecast_daily:
+        if daily_data := self.coordinator.data.daily:
             return {
-                ATTR_SMHI_THUNDER_PROBABILITY: self._forecast_daily[0]["thunder"],
+                ATTR_SMHI_THUNDER_PROBABILITY: daily_data[0]["thunder"],
             }
         return None
 
-    @Throttle(MIN_TIME_BETWEEN_UPDATES)
+    @callback
-    async def async_update(self) -> None:
+    def _handle_coordinator_update(self) -> None:
-        """Refresh the forecast data from SMHI weather API."""
+        """Handle updated data from the coordinator."""
-        try:
+        self.update_entity_data()
-            async with asyncio.timeout(TIMEOUT):
+        super()._handle_coordinator_update()
-                self._forecast_daily = await self._smhi_api.async_get_daily_forecast()
-                self._forecast_hourly = await self._smhi_api.async_get_hourly_forecast()
-                self._fail_count = 0
-        except (TimeoutError, SmhiForecastException):
-            _LOGGER.error("Failed to connect to SMHI API, retry in 5 minutes")
-            self._fail_count += 1
-            if self._fail_count < 3:
-                async_call_later(self.hass, RETRY_TIMEOUT, self.retry_update)
-                return
-
-        if self._forecast_daily:
-            self._attr_native_temperature = self._forecast_daily[0]["temperature"]
-            self._attr_humidity = self._forecast_daily[0]["humidity"]
-            self._attr_native_wind_speed = self._forecast_daily[0]["wind_speed"]
-            self._attr_wind_bearing = self._forecast_daily[0]["wind_direction"]
-            self._attr_native_visibility = self._forecast_daily[0]["visibility"]
-            self._attr_native_pressure = self._forecast_daily[0]["pressure"]
-            self._attr_native_wind_gust_speed = self._forecast_daily[0]["wind_gust"]
-            self._attr_cloud_coverage = self._forecast_daily[0]["total_cloud"]
-            self._attr_condition = CONDITION_MAP.get(self._forecast_daily[0]["symbol"])
-            if self._attr_condition == ATTR_CONDITION_SUNNY and not sun.is_up(
-                self.hass
-            ):
-                self._attr_condition = ATTR_CONDITION_CLEAR_NIGHT
-        await self.async_update_listeners(("daily", "hourly"))
-
-    async def retry_update(self, _: datetime) -> None:
-        """Retry refresh weather forecast."""
-        await self.async_update(no_throttle=True)
 
     def _get_forecast_data(
         self, forecast_data: list[SMHIForecast] | None
@@ -219,10 +188,10 @@ class SmhiWeather(SmhiWeatherBaseEntity, WeatherEntity):
 
         return data
 
-    async def async_forecast_daily(self) -> list[Forecast] | None:
+    def _async_forecast_daily(self) -> list[Forecast] | None:
         """Service to retrieve the daily forecast."""
-        return self._get_forecast_data(self._forecast_daily)
+        return self._get_forecast_data(self.coordinator.data.daily)
 
-    async def async_forecast_hourly(self) -> list[Forecast] | None:
+    def _async_forecast_hourly(self) -> list[Forecast] | None:
         """Service to retrieve the hourly forecast."""
-        return self._get_forecast_data(self._forecast_hourly)
+        return self._get_forecast_data(self.coordinator.data.hourly)

tests/components/smhi/test_weather.py

@@ -4,29 +4,27 @@ from datetime import datetime, timedelta
 from unittest.mock import patch
 
 from freezegun import freeze_time
+from freezegun.api import FrozenDateTimeFactory
 from pysmhi import SMHIForecast, SmhiForecastException
 from pysmhi.const import API_POINT_FORECAST
 import pytest
 from syrupy.assertion import SnapshotAssertion
 
-from homeassistant.components.smhi.const import ATTR_SMHI_THUNDER_PROBABILITY
+from homeassistant.components.smhi.weather import CONDITION_CLASSES
-from homeassistant.components.smhi.weather import CONDITION_CLASSES, RETRY_TIMEOUT
 from homeassistant.components.weather import (
     ATTR_CONDITION_CLEAR_NIGHT,
     ATTR_FORECAST_CONDITION,
-    ATTR_WEATHER_CLOUD_COVERAGE,
-    ATTR_WEATHER_HUMIDITY,
-    ATTR_WEATHER_PRESSURE,
-    ATTR_WEATHER_TEMPERATURE,
-    ATTR_WEATHER_VISIBILITY,
-    ATTR_WEATHER_WIND_BEARING,
     ATTR_WEATHER_WIND_GUST_SPEED,
-    ATTR_WEATHER_WIND_SPEED,
     ATTR_WEATHER_WIND_SPEED_UNIT,
     DOMAIN as WEATHER_DOMAIN,
     SERVICE_GET_FORECASTS,
 )
-from homeassistant.const import ATTR_ATTRIBUTION, STATE_UNKNOWN, UnitOfSpeed
+from homeassistant.const import (
+    ATTR_ATTRIBUTION,
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+    UnitOfSpeed,
+)
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 from homeassistant.util import dt as dt_util
@@ -104,33 +102,38 @@ async def test_clear_night(
     assert response == snapshot(name="clear-night_forecast")
 
 
-async def test_properties_no_data(hass: HomeAssistant) -> None:
+async def test_properties_no_data(
+    hass: HomeAssistant,
+    aioclient_mock: AiohttpClientMocker,
+    api_response: str,
+    freezer: FrozenDateTimeFactory,
+) -> None:
     """Test properties when no API data available."""
+    uri = API_POINT_FORECAST.format(
+        TEST_CONFIG["location"]["longitude"], TEST_CONFIG["location"]["latitude"]
+    )
+    aioclient_mock.get(uri, text=api_response)
+
     entry = MockConfigEntry(domain="smhi", title="test", data=TEST_CONFIG, version=3)
     entry.add_to_hass(hass)
 
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
     with patch(
-        "homeassistant.components.smhi.entity.SMHIPointForecast.async_get_daily_forecast",
+        "homeassistant.components.smhi.coordinator.SMHIPointForecast.async_get_daily_forecast",
         side_effect=SmhiForecastException("boom"),
     ):
-        await hass.config_entries.async_setup(entry.entry_id)
+        freezer.tick(timedelta(minutes=35))
+        async_fire_time_changed(hass)
         await hass.async_block_till_done()
 
     state = hass.states.get(ENTITY_ID)
 
     assert state
     assert state.name == "test"
-    assert state.state == STATE_UNKNOWN
+    assert state.state == STATE_UNAVAILABLE
     assert state.attributes[ATTR_ATTRIBUTION] == "Swedish weather institute (SMHI)"
-    assert ATTR_WEATHER_HUMIDITY not in state.attributes
-    assert ATTR_WEATHER_PRESSURE not in state.attributes
-    assert ATTR_WEATHER_TEMPERATURE not in state.attributes
-    assert ATTR_WEATHER_VISIBILITY not in state.attributes
-    assert ATTR_WEATHER_WIND_SPEED not in state.attributes
-    assert ATTR_WEATHER_WIND_BEARING not in state.attributes
-    assert ATTR_WEATHER_CLOUD_COVERAGE not in state.attributes
-    assert ATTR_SMHI_THUNDER_PROBABILITY not in state.attributes
-    assert ATTR_WEATHER_WIND_GUST_SPEED not in state.attributes
 
 
 async def test_properties_unknown_symbol(hass: HomeAssistant) -> None:
@@ -215,11 +218,11 @@ async def test_properties_unknown_symbol(hass: HomeAssistant) -> None:
 
     with (
         patch(
-            "homeassistant.components.smhi.entity.SMHIPointForecast.async_get_daily_forecast",
+            "homeassistant.components.smhi.coordinator.SMHIPointForecast.async_get_daily_forecast",
             return_value=testdata,
         ),
         patch(
-            "homeassistant.components.smhi.entity.SMHIPointForecast.async_get_hourly_forecast",
+            "homeassistant.components.smhi.coordinator.SMHIPointForecast.async_get_hourly_forecast",
             return_value=None,
         ),
     ):
@@ -246,55 +249,48 @@ async def test_properties_unknown_symbol(hass: HomeAssistant) -> None:
 
 @pytest.mark.parametrize("error", [SmhiForecastException(), TimeoutError()])
 async def test_refresh_weather_forecast_retry(
-    hass: HomeAssistant, error: Exception
+    hass: HomeAssistant,
+    error: Exception,
+    aioclient_mock: AiohttpClientMocker,
+    api_response: str,
+    freezer: FrozenDateTimeFactory,
 ) -> None:
     """Test the refresh weather forecast function."""
+    uri = API_POINT_FORECAST.format(
+        TEST_CONFIG["location"]["longitude"], TEST_CONFIG["location"]["latitude"]
+    )
+    aioclient_mock.get(uri, text=api_response)
+
     entry = MockConfigEntry(domain="smhi", title="test", data=TEST_CONFIG, version=3)
     entry.add_to_hass(hass)
-    now = dt_util.utcnow()
+
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
 
     with patch(
-        "homeassistant.components.smhi.entity.SMHIPointForecast.async_get_daily_forecast",
+        "homeassistant.components.smhi.coordinator.SMHIPointForecast.async_get_daily_forecast",
         side_effect=error,
     ) as mock_get_forecast:
-        await hass.config_entries.async_setup(entry.entry_id)
+        freezer.tick(timedelta(minutes=35))
+        async_fire_time_changed(hass)
         await hass.async_block_till_done()
 
         state = hass.states.get(ENTITY_ID)
 
         assert state
         assert state.name == "test"
-        assert state.state == STATE_UNKNOWN
+        assert state.state == STATE_UNAVAILABLE
         assert mock_get_forecast.call_count == 1
 
-        future = now + timedelta(seconds=RETRY_TIMEOUT + 1)
+        freezer.tick(timedelta(minutes=35))
-        async_fire_time_changed(hass, future)
+        async_fire_time_changed(hass)
         await hass.async_block_till_done()
 
         state = hass.states.get(ENTITY_ID)
         assert state
-        assert state.state == STATE_UNKNOWN
+        assert state.state == STATE_UNAVAILABLE
         assert mock_get_forecast.call_count == 2
 
-        future = future + timedelta(seconds=RETRY_TIMEOUT + 1)
-        async_fire_time_changed(hass, future)
-        await hass.async_block_till_done()
-
-        state = hass.states.get(ENTITY_ID)
-        assert state
-        assert state.state == STATE_UNKNOWN
-        assert mock_get_forecast.call_count == 3
-
-        future = future + timedelta(seconds=RETRY_TIMEOUT + 1)
-        async_fire_time_changed(hass, future)
-        await hass.async_block_till_done()
-
-        state = hass.states.get(ENTITY_ID)
-        assert state
-        assert state.state == STATE_UNKNOWN
-        # after three failed retries we stop retrying and go back to normal interval
-        assert mock_get_forecast.call_count == 3
-
 
 def test_condition_class() -> None:
     """Test condition class."""