Use DataUpdateCoordinator in NWS (#34372)

homeassistant/components/nws/__init__.py

@@ -3,7 +3,6 @@ import asyncio
 import datetime
 import logging
 
-import aiohttp
 from pynws import SimpleNWS
 import voluptuous as vol
 
@@ -12,10 +11,16 @@ from homeassistant.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 import homeassistant.helpers.config_validation as cv
-from homeassistant.helpers.dispatcher import async_dispatcher_send
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
-from homeassistant.helpers.event import async_track_time_interval
 
-from .const import CONF_STATION, DOMAIN
+from .const import (
+    CONF_STATION,
+    COORDINATOR_FORECAST,
+    COORDINATOR_FORECAST_HOURLY,
+    COORDINATOR_OBSERVATION,
+    DOMAIN,
+    NWS_DATA,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -46,11 +51,6 @@ def base_unique_id(latitude, longitude):
     return f"{latitude}_{longitude}"
 
 
-def signal_unique_id(latitude, longitude):
-    """Return unique id for signaling to entries in configuration from component."""
-    return f"{DOMAIN}_{base_unique_id(latitude,longitude)}"
-
-
 async def async_setup(hass: HomeAssistant, config: dict):
     """Set up the National Weather Service (NWS) component."""
     hass.data.setdefault(DOMAIN, {})
@@ -66,14 +66,45 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
 
     client_session = async_get_clientsession(hass)
 
-    nws_data = NwsData(hass, latitude, longitude, api_key, client_session)
+    # set_station only does IO when station is None
-    hass.data[DOMAIN][entry.entry_id] = nws_data
+    nws_data = SimpleNWS(latitude, longitude, api_key, client_session)
+    await nws_data.set_station(station)
 
-    # async_set_station only does IO when station is None
+    coordinator_observation = DataUpdateCoordinator(
-    await nws_data.async_set_station(station)
+        hass,
-    await nws_data.async_update()
+        _LOGGER,
+        name=f"NWS observation station {station}",
+        update_method=nws_data.update_observation,
+        update_interval=DEFAULT_SCAN_INTERVAL,
+    )
 
-    async_track_time_interval(hass, nws_data.async_update, DEFAULT_SCAN_INTERVAL)
+    coordinator_forecast = DataUpdateCoordinator(
+        hass,
+        _LOGGER,
+        name=f"NWS forecast station {station}",
+        update_method=nws_data.update_forecast,
+        update_interval=DEFAULT_SCAN_INTERVAL,
+    )
+
+    coordinator_forecast_hourly = DataUpdateCoordinator(
+        hass,
+        _LOGGER,
+        name=f"NWS forecast hourly station {station}",
+        update_method=nws_data.update_forecast_hourly,
+        update_interval=DEFAULT_SCAN_INTERVAL,
+    )
+
+    hass.data[DOMAIN][entry.entry_id] = {
+        NWS_DATA: nws_data,
+        COORDINATOR_OBSERVATION: coordinator_observation,
+        COORDINATOR_FORECAST: coordinator_forecast,
+        COORDINATOR_FORECAST_HOURLY: coordinator_forecast_hourly,
+    }
+
+    # Fetch initial data so we have data when entities subscribe
+    await coordinator_observation.async_refresh()
+    await coordinator_forecast.async_refresh()
+    await coordinator_forecast_hourly.async_refresh()
 
     for component in PLATFORMS:
         hass.async_create_task(
@@ -97,99 +128,3 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry):
         if len(hass.data[DOMAIN]) == 0:
             hass.data.pop(DOMAIN)
     return unload_ok
-
-
-class NwsData:
-    """Data class for National Weather Service integration."""
-
-    def __init__(self, hass, latitude, longitude, api_key, websession):
-        """Initialize the data."""
-        self.hass = hass
-        self.latitude = latitude
-        self.longitude = longitude
-        ha_api_key = f"{api_key} homeassistant"
-        self.nws = SimpleNWS(latitude, longitude, ha_api_key, websession)
-
-        self.update_observation_success = True
-        self.update_forecast_success = True
-        self.update_forecast_hourly_success = True
-
-    async def async_set_station(self, station):
-        """
-        Set to desired station.
-
-        If None, nearest station is used.
-        """
-        await self.nws.set_station(station)
-        _LOGGER.debug("Nearby station list: %s", self.nws.stations)
-
-    @property
-    def station(self):
-        """Return station name."""
-        return self.nws.station
-
-    @property
-    def observation(self):
-        """Return observation."""
-        return self.nws.observation
-
-    @property
-    def forecast(self):
-        """Return day+night forecast."""
-        return self.nws.forecast
-
-    @property
-    def forecast_hourly(self):
-        """Return hourly forecast."""
-        return self.nws.forecast_hourly
-
-    @staticmethod
-    async def _async_update_item(
-        update_call, update_type, station_name, previous_success
-    ):
-        """Update item and handle logging."""
-        try:
-            _LOGGER.debug("Updating %s for station %s", update_type, station_name)
-            await update_call()
-
-            if not previous_success:
-                _LOGGER.warning(
-                    "Success updating %s for station %s", update_type, station_name
-                )
-            success = True
-        except (aiohttp.ClientError, asyncio.TimeoutError) as err:
-            if previous_success:
-                _LOGGER.warning(
-                    "Error updating %s for station %s: %s",
-                    update_type,
-                    station_name,
-                    err,
-                )
-            success = False
-        return success
-
-    async def async_update(self, now=None):
-        """Update all data."""
-
-        self.update_observation_success = await self._async_update_item(
-            self.nws.update_observation,
-            "observation",
-            self.station,
-            self.update_observation_success,
-        )
-        self.update_forecast_success = await self._async_update_item(
-            self.nws.update_forecast,
-            "forecast",
-            self.station,
-            self.update_forecast_success,
-        )
-        self.update_forecast_hourly_success = await self._async_update_item(
-            self.nws.update_forecast_hourly,
-            "forecast_hourly",
-            self.station,
-            self.update_forecast_hourly_success,
-        )
-
-        async_dispatcher_send(
-            self.hass, signal_unique_id(self.latitude, self.longitude)
-        )

homeassistant/components/nws/config_flow.py

@@ -2,6 +2,7 @@
 import logging
 
 import aiohttp
+from pynws import SimpleNWS
 import voluptuous as vol
 
 from homeassistant import config_entries, core, exceptions
@@ -9,7 +10,7 @@ from homeassistant.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from . import NwsData, base_unique_id
+from . import base_unique_id
 from .const import CONF_STATION, DOMAIN  # pylint:disable=unused-import
 
 _LOGGER = logging.getLogger(__name__)
@@ -27,10 +28,10 @@ async def validate_input(hass: core.HomeAssistant, data):
 
     client_session = async_get_clientsession(hass)
     ha_api_key = f"{api_key} homeassistant"
-    nws = NwsData(hass, latitude, longitude, ha_api_key, client_session)
+    nws = SimpleNWS(latitude, longitude, ha_api_key, client_session)
 
     try:
-        await nws.async_set_station(station)
+        await nws.set_station(station)
     except aiohttp.ClientError as err:
         _LOGGER.error("Could not connect: %s", err)
         raise CannotConnect

homeassistant/components/nws/const.py

@@ -55,3 +55,8 @@ CONDITION_CLASSES = {
 
 DAYNIGHT = "daynight"
 HOURLY = "hourly"
+
+NWS_DATA = "nws data"
+COORDINATOR_OBSERVATION = "coordinator_observation"
+COORDINATOR_FORECAST = "coordinator_forecast"
+COORDINATOR_FORECAST_HOURLY = "coordinator_forecast_hourly"

homeassistant/components/nws/weather.py

@@ -10,6 +10,8 @@ from homeassistant.components.weather import (
     WeatherEntity,
 )
 from homeassistant.const import (
+    CONF_LATITUDE,
+    CONF_LONGITUDE,
     LENGTH_KILOMETERS,
     LENGTH_METERS,
     LENGTH_MILES,
@@ -20,22 +22,25 @@ from homeassistant.const import (
     TEMP_FAHRENHEIT,
 )
 from homeassistant.core import callback
-from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.typing import ConfigType, HomeAssistantType
 from homeassistant.util.distance import convert as convert_distance
 from homeassistant.util.pressure import convert as convert_pressure
 from homeassistant.util.temperature import convert as convert_temperature
 
-from . import base_unique_id, signal_unique_id
+from . import base_unique_id
 from .const import (
     ATTR_FORECAST_DAYTIME,
     ATTR_FORECAST_DETAILED_DESCRIPTION,
     ATTR_FORECAST_PRECIP_PROB,
     ATTRIBUTION,
     CONDITION_CLASSES,
+    COORDINATOR_FORECAST,
+    COORDINATOR_FORECAST_HOURLY,
+    COORDINATOR_OBSERVATION,
     DAYNIGHT,
     DOMAIN,
     HOURLY,
+    NWS_DATA,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -73,11 +78,12 @@ async def async_setup_entry(
     hass: HomeAssistantType, entry: ConfigType, async_add_entities
 ) -> None:
     """Set up the NWS weather platform."""
-    nws_data = hass.data[DOMAIN][entry.entry_id]
+    hass_data = hass.data[DOMAIN][entry.entry_id]
+
     async_add_entities(
         [
-            NWSWeather(nws_data, DAYNIGHT, hass.config.units),
+            NWSWeather(entry.data, hass_data, DAYNIGHT, hass.config.units),
-            NWSWeather(nws_data, HOURLY, hass.config.units),
+            NWSWeather(entry.data, hass_data, HOURLY, hass.config.units),
         ],
         False,
     )
@@ -86,12 +92,17 @@ async def async_setup_entry(
 class NWSWeather(WeatherEntity):
     """Representation of a weather condition."""
 
-    def __init__(self, nws, mode, units):
+    def __init__(self, entry_data, hass_data, mode, units):
         """Initialise the platform with a data instance and station name."""
-        self.nws = nws
+        self.nws = hass_data[NWS_DATA]
-        self.station = nws.station
+        self.latitude = entry_data[CONF_LATITUDE]
-        self.latitude = nws.latitude
+        self.longitude = entry_data[CONF_LONGITUDE]
-        self.longitude = nws.longitude
+        self.coordinator_observation = hass_data[COORDINATOR_OBSERVATION]
+        if mode == DAYNIGHT:
+            self.coordinator_forecast = hass_data[COORDINATOR_FORECAST]
+        else:
+            self.coordinator_forecast = hass_data[COORDINATOR_FORECAST_HOURLY]
+        self.station = self.nws.station
 
         self.is_metric = units.is_metric
         self.mode = mode
@@ -102,11 +113,10 @@ class NWSWeather(WeatherEntity):
     async def async_added_to_hass(self) -> None:
         """Set up a listener and load data."""
         self.async_on_remove(
-            async_dispatcher_connect(
+            self.coordinator_observation.async_add_listener(self._update_callback)
-                self.hass,
+        )
-                signal_unique_id(self.latitude, self.longitude),
+        self.async_on_remove(
-                self._update_callback,
+            self.coordinator_forecast.async_add_listener(self._update_callback)
-            )
         )
         self._update_callback()
 
@@ -275,11 +285,7 @@ class NWSWeather(WeatherEntity):
     @property
     def available(self):
         """Return if state is available."""
-        if self.mode == DAYNIGHT:
-            return (
-                self.nws.update_observation_success and self.nws.update_forecast_success
-            )
         return (
-            self.nws.update_observation_success
+            self.coordinator_observation.last_update_success
-            and self.nws.update_forecast_hourly_success
+            and self.coordinator_forecast.last_update_success
         )

tests/components/nws/conftest.py

@@ -22,3 +22,14 @@ def mock_simple_nws():
         instance.forecast = DEFAULT_FORECAST
         instance.forecast_hourly = DEFAULT_FORECAST
         yield mock_nws
+
+
+@pytest.fixture()
+def mock_simple_nws_config():
+    """Mock pynws SimpleNWS with default values in config_flow."""
+    with patch("homeassistant.components.nws.config_flow.SimpleNWS") as mock_nws:
+        instance = mock_nws.return_value
+        instance.set_station.return_value = mock_coro()
+        instance.station = "ABC"
+        instance.stations = ["ABC"]
+        yield mock_nws

tests/components/nws/test_config_flow.py

@@ -6,7 +6,7 @@ from homeassistant import config_entries, setup
 from homeassistant.components.nws.const import DOMAIN
 
 
-async def test_form(hass, mock_simple_nws):
+async def test_form(hass, mock_simple_nws_config):
     """Test we get the form."""
     hass.config.latitude = 35
     hass.config.longitude = -90
@@ -40,9 +40,9 @@ async def test_form(hass, mock_simple_nws):
     assert len(mock_setup_entry.mock_calls) == 1
 
 
-async def test_form_cannot_connect(hass, mock_simple_nws):
+async def test_form_cannot_connect(hass, mock_simple_nws_config):
     """Test we handle cannot connect error."""
-    mock_instance = mock_simple_nws.return_value
+    mock_instance = mock_simple_nws_config.return_value
     mock_instance.set_station.side_effect = aiohttp.ClientError
 
     result = await hass.config_entries.flow.async_init(
@@ -57,9 +57,9 @@ async def test_form_cannot_connect(hass, mock_simple_nws):
     assert result2["errors"] == {"base": "cannot_connect"}
 
 
-async def test_form_unknown_error(hass, mock_simple_nws):
+async def test_form_unknown_error(hass, mock_simple_nws_config):
     """Test we handle unknown error."""
-    mock_instance = mock_simple_nws.return_value
+    mock_instance = mock_simple_nws_config.return_value
     mock_instance.set_station.side_effect = ValueError
 
     result = await hass.config_entries.flow.async_init(
@@ -74,7 +74,7 @@ async def test_form_unknown_error(hass, mock_simple_nws):
     assert result2["errors"] == {"base": "unknown"}
 
 
-async def test_form_already_configured(hass, mock_simple_nws):
+async def test_form_already_configured(hass, mock_simple_nws_config):
     """Test we handle duplicate entries."""
     result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": config_entries.SOURCE_USER}

tests/components/nws/test_weather.py

@@ -29,7 +29,7 @@ from tests.components.nws.const import (
     ],
 )
 async def test_imperial_metric(
-    hass, units, result_observation, result_forecast, mock_simple_nws, caplog
+    hass, units, result_observation, result_forecast, mock_simple_nws
 ):
     """Test with imperial and metric units."""
     hass.config.units = units
@@ -64,9 +64,6 @@ async def test_imperial_metric(
     for key, value in result_forecast.items():
         assert forecast[0].get(key) == value
 
-    assert "Error updating observation" not in caplog.text
-    assert "Success updating observation" not in caplog.text
-
 
 async def test_none_values(hass, mock_simple_nws):
     """Test with none values in observation and forecast dicts."""
@@ -128,7 +125,7 @@ async def test_error_station(hass, mock_simple_nws):
     assert hass.states.get("weather.abc_daynight") is None
 
 
-async def test_error_observation(hass, mock_simple_nws, caplog):
+async def test_error_observation(hass, mock_simple_nws):
     """Test error during update observation."""
     instance = mock_simple_nws.return_value
     instance.update_observation.side_effect = aiohttp.ClientError
@@ -148,10 +145,6 @@ async def test_error_observation(hass, mock_simple_nws, caplog):
     assert state
     assert state.state == "unavailable"
 
-    assert "Error updating observation for station ABC" in caplog.text
-    assert "Success updating observation for station ABC" not in caplog.text
-    caplog.clear()
-
     instance.update_observation.side_effect = None
 
     future_time = dt_util.utcnow() + timedelta(minutes=15)
@@ -168,11 +161,8 @@ async def test_error_observation(hass, mock_simple_nws, caplog):
     assert state
     assert state.state == "sunny"
 
-    assert "Error updating observation for station ABC" not in caplog.text
-    assert "Success updating observation for station ABC" in caplog.text
 
-
+async def test_error_forecast(hass, mock_simple_nws):
-async def test_error_forecast(hass, caplog, mock_simple_nws):
     """Test error during update forecast."""
     instance = mock_simple_nws.return_value
     instance.update_forecast.side_effect = aiohttp.ClientError
@@ -188,10 +178,6 @@ async def test_error_forecast(hass, caplog, mock_simple_nws):
     assert state
     assert state.state == "unavailable"
 
-    assert "Error updating forecast for station ABC" in caplog.text
-    assert "Success updating forecast for station ABC" not in caplog.text
-    caplog.clear()
-
     instance.update_forecast.side_effect = None
 
     future_time = dt_util.utcnow() + timedelta(minutes=15)
@@ -204,11 +190,8 @@ async def test_error_forecast(hass, caplog, mock_simple_nws):
     assert state
     assert state.state == "sunny"
 
-    assert "Error updating forecast for station ABC" not in caplog.text
-    assert "Success updating forecast for station ABC" in caplog.text
 
-
+async def test_error_forecast_hourly(hass, mock_simple_nws):
-async def test_error_forecast_hourly(hass, caplog, mock_simple_nws):
     """Test error during update forecast hourly."""
     instance = mock_simple_nws.return_value
     instance.update_forecast_hourly.side_effect = aiohttp.ClientError
@@ -224,10 +207,6 @@ async def test_error_forecast_hourly(hass, caplog, mock_simple_nws):
 
     instance.update_forecast_hourly.assert_called_once()
 
-    assert "Error updating forecast_hourly for station ABC" in caplog.text
-    assert "Success updating forecast_hourly for station ABC" not in caplog.text
-    caplog.clear()
-
     instance.update_forecast_hourly.side_effect = None
 
     future_time = dt_util.utcnow() + timedelta(minutes=15)
@@ -239,6 +218,3 @@ async def test_error_forecast_hourly(hass, caplog, mock_simple_nws):
     state = hass.states.get("weather.abc_hourly")
     assert state
     assert state.state == "sunny"
-
-    assert "Error updating forecast_hourly for station ABC" not in caplog.text
-    assert "Success updating forecast_hourly for station ABC" in caplog.text