Modernize tomorrowio weather (#98466)

* Modernize tomorrowio weather

* Add test snapshot

* Update snapshots

* Address review comments

* Improve test coverage

homeassistant/components/tomorrowio/weather.py

@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 from datetime import datetime
-from typing import Any
 
 from pytomorrowio.const import DAILY, FORECASTS, HOURLY, NOWCAST, WeatherCode
 
@@ -15,7 +14,10 @@ from homeassistant.components.weather import (
     ATTR_FORECAST_PRECIPITATION_PROBABILITY,
     ATTR_FORECAST_TIME,
     ATTR_FORECAST_WIND_BEARING,
+    DOMAIN as WEATHER_DOMAIN,
+    Forecast,
     WeatherEntity,
+    WeatherEntityFeature,
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
@@ -27,7 +29,8 @@ from homeassistant.const import (
     UnitOfSpeed,
     UnitOfTemperature,
 )
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.sun import is_up
 from homeassistant.util import dt as dt_util
@@ -63,14 +66,30 @@ async def async_setup_entry(
 ) -> None:
     """Set up a config entry."""
     coordinator = hass.data[DOMAIN][config_entry.data[CONF_API_KEY]]
+    entity_registry = er.async_get(hass)
 
-    entities = [
+    entities = [TomorrowioWeatherEntity(config_entry, coordinator, 4, DAILY)]
+
+    # Add hourly and nowcast entities to legacy config entries
+    for forecast_type in (HOURLY, NOWCAST):
+        if not entity_registry.async_get_entity_id(
+            WEATHER_DOMAIN,
+            DOMAIN,
+            _calculate_unique_id(config_entry.unique_id, forecast_type),
+        ):
+            continue
+        entities.append(
             TomorrowioWeatherEntity(config_entry, coordinator, 4, forecast_type)
-        for forecast_type in (DAILY, HOURLY, NOWCAST)
+        )
-    ]
+
     async_add_entities(entities)
 
 
+def _calculate_unique_id(config_entry_unique_id: str | None, forecast_type: str) -> str:
+    """Calculate unique ID."""
+    return f"{config_entry_unique_id}_{forecast_type}"
+
+
 class TomorrowioWeatherEntity(TomorrowioEntity, WeatherEntity):
     """Entity that talks to Tomorrow.io v4 API to retrieve weather data."""
 
@@ -79,6 +98,9 @@ class TomorrowioWeatherEntity(TomorrowioEntity, WeatherEntity):
     _attr_native_temperature_unit = UnitOfTemperature.CELSIUS
     _attr_native_visibility_unit = UnitOfLength.KILOMETERS
     _attr_native_wind_speed_unit = UnitOfSpeed.METERS_PER_SECOND
+    _attr_supported_features = (
+        WeatherEntityFeature.FORECAST_DAILY | WeatherEntityFeature.FORECAST_HOURLY
+    )
 
     def __init__(
         self,
@@ -94,7 +116,18 @@ class TomorrowioWeatherEntity(TomorrowioEntity, WeatherEntity):
             forecast_type == DEFAULT_FORECAST_TYPE
         )
         self._attr_name = f"{config_entry.data[CONF_NAME]} - {forecast_type.title()}"
-        self._attr_unique_id = f"{config_entry.unique_id}_{forecast_type}"
+        self._attr_unique_id = _calculate_unique_id(
+            config_entry.unique_id, forecast_type
+        )
+
+    @callback
+    def _handle_coordinator_update(self) -> None:
+        """Handle updated data from the coordinator."""
+        super()._handle_coordinator_update()
+        assert self.platform.config_entry
+        self.platform.config_entry.async_create_task(
+            self.hass, self.async_update_listeners(("daily", "hourly"))
+        )
 
     def _forecast_dict(
         self,
@@ -102,12 +135,12 @@ class TomorrowioWeatherEntity(TomorrowioEntity, WeatherEntity):
         use_datetime: bool,
         condition: int,
         precipitation: float | None,
-        precipitation_probability: float | None,
+        precipitation_probability: int | None,
         temp: float | None,
         temp_low: float | None,
         wind_direction: float | None,
         wind_speed: float | None,
-    ) -> dict[str, Any]:
+    ) -> Forecast:
         """Return formatted Forecast dict from Tomorrow.io forecast data."""
         if use_datetime:
             translated_condition = self._translate_condition(
@@ -116,7 +149,7 @@ class TomorrowioWeatherEntity(TomorrowioEntity, WeatherEntity):
         else:
             translated_condition = self._translate_condition(condition, True)
 
-        data = {
+        return {
             ATTR_FORECAST_TIME: forecast_dt.isoformat(),
             ATTR_FORECAST_CONDITION: translated_condition,
             ATTR_FORECAST_NATIVE_PRECIPITATION: precipitation,
@@ -127,8 +160,6 @@ class TomorrowioWeatherEntity(TomorrowioEntity, WeatherEntity):
             ATTR_FORECAST_NATIVE_WIND_SPEED: wind_speed,
         }
 
-        return {k: v for k, v in data.items() if v is not None}
-
     @staticmethod
     def _translate_condition(
         condition: int | None, sun_is_up: bool = True
@@ -187,20 +218,19 @@ class TomorrowioWeatherEntity(TomorrowioEntity, WeatherEntity):
         """Return the raw visibility."""
         return self._get_current_property(TMRW_ATTR_VISIBILITY)
 
-    @property
+    def _forecast(self, forecast_type: str) -> list[Forecast] | None:
-    def forecast(self):
         """Return the forecast."""
         # Check if forecasts are available
         raw_forecasts = (
             self.coordinator.data.get(self._config_entry.entry_id, {})
             .get(FORECASTS, {})
-            .get(self.forecast_type)
+            .get(forecast_type)
         )
         if not raw_forecasts:
             return None
 
-        forecasts = []
+        forecasts: list[Forecast] = []
-        max_forecasts = MAX_FORECASTS[self.forecast_type]
+        max_forecasts = MAX_FORECASTS[forecast_type]
         forecast_count = 0
 
         # Convert utcnow to local to be compatible with tests
@@ -212,7 +242,7 @@ class TomorrowioWeatherEntity(TomorrowioEntity, WeatherEntity):
             forecast_dt = dt_util.parse_datetime(forecast[TMRW_ATTR_TIMESTAMP])
 
             # Throw out past data
-            if dt_util.as_local(forecast_dt).date() < today:
+            if forecast_dt is None or dt_util.as_local(forecast_dt).date() < today:
                 continue
 
             values = forecast["values"]
@@ -222,18 +252,23 @@ class TomorrowioWeatherEntity(TomorrowioEntity, WeatherEntity):
             precipitation = values.get(TMRW_ATTR_PRECIPITATION)
             precipitation_probability = values.get(TMRW_ATTR_PRECIPITATION_PROBABILITY)
 
+            try:
+                precipitation_probability = round(precipitation_probability)
+            except TypeError:
+                precipitation_probability = None
+
             temp = values.get(TMRW_ATTR_TEMPERATURE_HIGH)
             temp_low = None
 
             wind_direction = values.get(TMRW_ATTR_WIND_DIRECTION)
             wind_speed = values.get(TMRW_ATTR_WIND_SPEED)
 
-            if self.forecast_type == DAILY:
+            if forecast_type == DAILY:
                 use_datetime = False
                 temp_low = values.get(TMRW_ATTR_TEMPERATURE_LOW)
                 if precipitation:
                     precipitation = precipitation * 24
-            elif self.forecast_type == NOWCAST:
+            elif forecast_type == NOWCAST:
                 # Precipitation is forecasted in CONF_TIMESTEP increments but in a
                 # per hour rate, so value needs to be converted to an amount.
                 if precipitation:
@@ -260,3 +295,16 @@ class TomorrowioWeatherEntity(TomorrowioEntity, WeatherEntity):
                 break
 
         return forecasts
+
+    @property
+    def forecast(self) -> list[Forecast] | None:
+        """Return the forecast array."""
+        return self._forecast(self.forecast_type)
+
+    async def async_forecast_daily(self) -> list[Forecast] | None:
+        """Return the daily forecast in native units."""
+        return self._forecast(DAILY)
+
+    async def async_forecast_hourly(self) -> list[Forecast] | None:
+        """Return the hourly forecast in native units."""
+        return self._forecast(HOURLY)

tests/components/tomorrowio/test_weather.py

@@ -1,10 +1,13 @@
 """Tests for Tomorrow.io weather entity."""
 from __future__ import annotations
 
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import Any
 
 from freezegun import freeze_time
+from freezegun.api import FrozenDateTimeFactory
+import pytest
+from syrupy.assertion import SnapshotAssertion
 
 from homeassistant.components.tomorrowio.config_flow import (
     _get_config_schema,
@@ -41,16 +44,19 @@ from homeassistant.components.weather import (
     ATTR_WEATHER_WIND_SPEED,
     ATTR_WEATHER_WIND_SPEED_UNIT,
     DOMAIN as WEATHER_DOMAIN,
+    SERVICE_GET_FORECAST,
 )
 from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY, SOURCE_USER
 from homeassistant.const import ATTR_ATTRIBUTION, ATTR_FRIENDLY_NAME, CONF_NAME
 from homeassistant.core import HomeAssistant, State, callback
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_registry import async_get
 from homeassistant.util import dt as dt_util
 
 from .const import API_V4_ENTRY_DATA
 
 from tests.common import MockConfigEntry, async_fire_time_changed
+from tests.typing import WebSocketGenerator
 
 
 @callback
@@ -65,11 +71,8 @@ def _enable_entity(hass: HomeAssistant, entity_name: str) -> None:
     assert updated_entry.disabled is False
 
 
-async def _setup(hass: HomeAssistant, config: dict[str, Any]) -> State:
+async def _setup_config_entry(hass: HomeAssistant, config: dict[str, Any]) -> State:
     """Set up entry and return entity state."""
-    with freeze_time(
-        datetime(2021, 3, 6, 23, 59, 59, tzinfo=dt_util.UTC)
-    ) as frozen_time:
     data = _get_config_schema(hass, SOURCE_USER)(config)
     data[CONF_NAME] = DEFAULT_NAME
     config_entry = MockConfigEntry(
@@ -82,6 +85,33 @@ async def _setup(hass: HomeAssistant, config: dict[str, Any]) -> State:
     config_entry.add_to_hass(hass)
     assert await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
+
+
+async def _setup(hass: HomeAssistant, config: dict[str, Any]) -> State:
+    """Set up entry and return entity state."""
+    with freeze_time(datetime(2021, 3, 6, 23, 59, 59, tzinfo=dt_util.UTC)):
+        await _setup_config_entry(hass, config)
+
+    return hass.states.get("weather.tomorrow_io_daily")
+
+
+async def _setup_legacy(hass: HomeAssistant, config: dict[str, Any]) -> State:
+    """Set up entry and return entity state."""
+    registry = er.async_get(hass)
+    data = _get_config_schema(hass, SOURCE_USER)(config)
+    for entity_name in ("hourly", "nowcast"):
+        registry.async_get_or_create(
+            WEATHER_DOMAIN,
+            DOMAIN,
+            f"{_get_unique_id(hass, data)}_{entity_name}",
+            disabled_by=er.RegistryEntryDisabler.INTEGRATION,
+            suggested_object_id=f"tomorrow_io_{entity_name}",
+        )
+
+    with freeze_time(
+        datetime(2021, 3, 6, 23, 59, 59, tzinfo=dt_util.UTC)
+    ) as frozen_time:
+        await _setup_config_entry(hass, config)
         for entity_name in ("hourly", "nowcast"):
             _enable_entity(hass, f"weather.tomorrow_io_{entity_name}")
         await hass.async_block_till_done()
@@ -94,6 +124,33 @@ async def _setup(hass: HomeAssistant, config: dict[str, Any]) -> State:
     return hass.states.get("weather.tomorrow_io_daily")
 
 
+async def test_new_config_entry(hass: HomeAssistant) -> None:
+    """Test the expected entities are created."""
+    registry = er.async_get(hass)
+    await _setup(hass, API_V4_ENTRY_DATA)
+    assert len(hass.states.async_entity_ids("weather")) == 1
+
+    entry = hass.config_entries.async_entries()[0]
+    assert len(er.async_entries_for_config_entry(registry, entry.entry_id)) == 28
+
+
+async def test_legacy_config_entry(hass: HomeAssistant) -> None:
+    """Test the expected entities are created."""
+    registry = er.async_get(hass)
+    data = _get_config_schema(hass, SOURCE_USER)(API_V4_ENTRY_DATA)
+    for entity_name in ("hourly", "nowcast"):
+        registry.async_get_or_create(
+            WEATHER_DOMAIN,
+            DOMAIN,
+            f"{_get_unique_id(hass, data)}_{entity_name}",
+        )
+    await _setup(hass, API_V4_ENTRY_DATA)
+    assert len(hass.states.async_entity_ids("weather")) == 3
+
+    entry = hass.config_entries.async_entries()[0]
+    assert len(er.async_entries_for_config_entry(registry, entry.entry_id)) == 30
+
+
 async def test_v4_weather(hass: HomeAssistant) -> None:
     """Test v4 weather data."""
     weather_state = await _setup(hass, API_V4_ENTRY_DATA)
@@ -123,3 +180,136 @@ async def test_v4_weather(hass: HomeAssistant) -> None:
     assert weather_state.attributes[ATTR_WEATHER_WIND_BEARING] == 315.14
     assert weather_state.attributes[ATTR_WEATHER_WIND_SPEED] == 33.59  # 9.33 m/s ->km/h
     assert weather_state.attributes[ATTR_WEATHER_WIND_SPEED_UNIT] == "km/h"
+
+
+async def test_v4_weather_legacy_entities(hass: HomeAssistant) -> None:
+    """Test v4 weather data."""
+    weather_state = await _setup_legacy(hass, API_V4_ENTRY_DATA)
+    assert weather_state.state == ATTR_CONDITION_SUNNY
+    assert weather_state.attributes[ATTR_ATTRIBUTION] == ATTRIBUTION
+    assert len(weather_state.attributes[ATTR_FORECAST]) == 14
+    assert weather_state.attributes[ATTR_FORECAST][0] == {
+        ATTR_FORECAST_CONDITION: ATTR_CONDITION_SUNNY,
+        ATTR_FORECAST_TIME: "2021-03-07T11:00:00+00:00",
+        ATTR_FORECAST_PRECIPITATION: 0,
+        ATTR_FORECAST_PRECIPITATION_PROBABILITY: 0,
+        ATTR_FORECAST_TEMP: 45.9,
+        ATTR_FORECAST_TEMP_LOW: 26.1,
+        ATTR_FORECAST_WIND_BEARING: 239.6,
+        ATTR_FORECAST_WIND_SPEED: 34.16,  # 9.49 m/s -> km/h
+    }
+    assert weather_state.attributes[ATTR_FRIENDLY_NAME] == "Tomorrow.io - Daily"
+    assert weather_state.attributes[ATTR_WEATHER_HUMIDITY] == 23
+    assert weather_state.attributes[ATTR_WEATHER_OZONE] == 46.53
+    assert weather_state.attributes[ATTR_WEATHER_PRECIPITATION_UNIT] == "mm"
+    assert weather_state.attributes[ATTR_WEATHER_PRESSURE] == 30.35
+    assert weather_state.attributes[ATTR_WEATHER_PRESSURE_UNIT] == "hPa"
+    assert weather_state.attributes[ATTR_WEATHER_TEMPERATURE] == 44.1
+    assert weather_state.attributes[ATTR_WEATHER_TEMPERATURE_UNIT] == "°C"
+    assert weather_state.attributes[ATTR_WEATHER_VISIBILITY] == 8.15
+    assert weather_state.attributes[ATTR_WEATHER_VISIBILITY_UNIT] == "km"
+    assert weather_state.attributes[ATTR_WEATHER_WIND_BEARING] == 315.14
+    assert weather_state.attributes[ATTR_WEATHER_WIND_SPEED] == 33.59  # 9.33 m/s ->km/h
+    assert weather_state.attributes[ATTR_WEATHER_WIND_SPEED_UNIT] == "km/h"
+
+
+@freeze_time(datetime(2021, 3, 6, 23, 59, 59, tzinfo=dt_util.UTC))
+async def test_v4_forecast_service(
+    hass: HomeAssistant,
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test multiple forecast."""
+    weather_state = await _setup(hass, API_V4_ENTRY_DATA)
+    entity_id = weather_state.entity_id
+
+    for forecast_type in ("daily", "hourly"):
+        response = await hass.services.async_call(
+            WEATHER_DOMAIN,
+            SERVICE_GET_FORECAST,
+            {
+                "entity_id": entity_id,
+                "type": forecast_type,
+            },
+            blocking=True,
+            return_response=True,
+        )
+        assert response["forecast"] != []
+        assert response == snapshot
+
+
+async def test_v4_bad_forecast(
+    hass: HomeAssistant,
+    freezer: FrozenDateTimeFactory,
+    tomorrowio_config_entry_update,
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test bad forecast data."""
+    freezer.move_to(datetime(2021, 3, 6, 23, 59, 59, tzinfo=dt_util.UTC))
+
+    weather_state = await _setup(hass, API_V4_ENTRY_DATA)
+    entity_id = weather_state.entity_id
+    hourly_forecast = tomorrowio_config_entry_update.return_value["forecasts"]["hourly"]
+    hourly_forecast[0]["values"]["precipitationProbability"] = "blah"
+
+    # Trigger data refetch
+    freezer.tick(timedelta(minutes=32) + timedelta(seconds=1))
+    await hass.async_block_till_done()
+
+    response = await hass.services.async_call(
+        WEATHER_DOMAIN,
+        SERVICE_GET_FORECAST,
+        {
+            "entity_id": entity_id,
+            "type": "hourly",
+        },
+        blocking=True,
+        return_response=True,
+    )
+    assert response["forecast"][0]["precipitation_probability"] is None
+
+
+@pytest.mark.parametrize("forecast_type", ["daily", "hourly"])
+async def test_forecast_subscription(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    freezer: FrozenDateTimeFactory,
+    snapshot: SnapshotAssertion,
+    forecast_type: str,
+) -> None:
+    """Test multiple forecast."""
+    client = await hass_ws_client(hass)
+    freezer.move_to(datetime(2021, 3, 6, 23, 59, 59, tzinfo=dt_util.UTC))
+
+    weather_state = await _setup(hass, API_V4_ENTRY_DATA)
+    entity_id = weather_state.entity_id
+
+    await client.send_json_auto_id(
+        {
+            "type": "weather/subscribe_forecast",
+            "forecast_type": forecast_type,
+            "entity_id": entity_id,
+        }
+    )
+    msg = await client.receive_json()
+    assert msg["success"]
+    assert msg["result"] is None
+    subscription_id = msg["id"]
+
+    msg = await client.receive_json()
+    assert msg["id"] == subscription_id
+    assert msg["type"] == "event"
+    forecast1 = msg["event"]["forecast"]
+
+    assert forecast1 != []
+    assert forecast1 == snapshot
+
+    freezer.tick(timedelta(minutes=32) + timedelta(seconds=1))
+    await hass.async_block_till_done()
+    msg = await client.receive_json()
+
+    assert msg["id"] == subscription_id
+    assert msg["type"] == "event"
+    forecast2 = msg["event"]["forecast"]
+
+    assert forecast2 != []
+    assert forecast2 == snapshot