Make tomorrowio API rate limit handling more robust (#70412)

* Use the max request limit when setting tomorrowio update interval

* tests

* reduce lines

* simplify

* refactor

* Make Coordinator.async_setup_entry more efficient at determining when to refresh data and schedule refresh

* clean up

* clean up

* Remove unnecessary type definition

* typo

* fix logic ot be more deterministic

* Another fix

* Comment

* Reduce wasted API calls by doing partial updates when new entries get added with a new key

* Simplify and use asyncio event so that config entries only load after initial coordinator refresh

* Remove commented out piece

* Comment

* Remove unnecessary variable

* More cleanup

* Make future merge easier

* remove dupe

* switch order

* add comment

* Remove unnecessary error handling

* make code easier to read

* review feedback for code

* Fix logic

* Update test based on review

* Tweak comments

* reset mock so asertions are more clear

* Remove update interval check

homeassistant/components/tomorrowio/__init__.py

@@ -1,6 +1,7 @@
 """The Tomorrow.io integration."""
 from __future__ import annotations
 
+import asyncio
 from datetime import timedelta
 import logging
 from math import ceil
@@ -23,7 +24,6 @@ from homeassistant.const import (
     CONF_LATITUDE,
     CONF_LOCATION,
     CONF_LONGITUDE,
-    CONF_NAME,
 )
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr, entity_registry as er
@@ -40,7 +40,6 @@ from .const import (
     CONF_TIMESTEP,
     DOMAIN,
     INTEGRATION_NAME,
-    MAX_REQUESTS_PER_DAY,
     TMRW_ATTR_CARBON_MONOXIDE,
     TMRW_ATTR_CHINA_AQI,
     TMRW_ATTR_CHINA_HEALTH_CONCERN,
@@ -85,36 +84,33 @@ PLATFORMS = [SENSOR_DOMAIN, WEATHER_DOMAIN]
 
 
 @callback
-def async_set_update_interval(
+def async_get_entries_by_api_key(
-    hass: HomeAssistant, current_entry: ConfigEntry
+    hass: HomeAssistant, api_key: str, exclude_entry: ConfigEntry | None = None
-) -> timedelta:
+) -> list[ConfigEntry]:
-    """Recalculate update_interval based on existing Tomorrow.io instances and update them."""
+    """Get all entries for a given API key."""
-    api_calls = 2
+    return [
-    # We check how many Tomorrow.io configured instances are using the same API key and
+        entry
-    # calculate interval to not exceed allowed numbers of requests. Divide 90% of
-    # MAX_REQUESTS_PER_DAY by the number of API calls because we want a buffer in the
-    # number of API calls left at the end of the day.
-    other_instance_entry_ids = [
-        entry.entry_id
         for entry in hass.config_entries.async_entries(DOMAIN)
-        if entry.entry_id != current_entry.entry_id
+        if entry.data[CONF_API_KEY] == api_key
-        and entry.data[CONF_API_KEY] == current_entry.data[CONF_API_KEY]
+        and (exclude_entry is None or exclude_entry != entry)
     ]
 
-    interval = timedelta(
+
-        minutes=(
+@callback
-            ceil(
+def async_set_update_interval(
-                (24 * 60 * (len(other_instance_entry_ids) + 1) * api_calls)
+    hass: HomeAssistant, api: TomorrowioV4, exclude_entry: ConfigEntry | None = None
-                / (MAX_REQUESTS_PER_DAY * 0.9)
+) -> timedelta:
-            )
+    """Calculate update_interval."""
-        )
+    # We check how many Tomorrow.io configured instances are using the same API key and
+    # calculate interval to not exceed allowed numbers of requests. Divide 90% of
+    # max_requests by the number of API calls because we want a buffer in the
+    # number of API calls left at the end of the day.
+    entries = async_get_entries_by_api_key(hass, api.api_key, exclude_entry)
+    minutes = ceil(
+        (24 * 60 * len(entries) * api.num_api_requests)
+        / (api.max_requests_per_day * 0.9)
     )
-
+    return timedelta(minutes=minutes)
-    for entry_id in other_instance_entry_ids:
-        if entry_id in hass.data[DOMAIN]:
-            hass.data[DOMAIN][entry_id].update_interval = interval
-
-    return interval
 
 
 @callback
@@ -197,24 +193,18 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     if entry.source == SOURCE_IMPORT and "old_config_entry_id" in entry.data:
         async_migrate_entry_from_climacell(hass, dev_reg, entry, device)
 
-    api = TomorrowioV4(
+    api_key = entry.data[CONF_API_KEY]
-        entry.data[CONF_API_KEY],
+    # If coordinator already exists for this API key, we'll use that, otherwise
-        entry.data[CONF_LOCATION][CONF_LATITUDE],
+    # we have to create a new one
-        entry.data[CONF_LOCATION][CONF_LONGITUDE],
+    if not (coordinator := hass.data[DOMAIN].get(api_key)):
-        unit_system="metric",
+        session = async_get_clientsession(hass)
-        session=async_get_clientsession(hass),
+        # we will not use the class's lat and long so we can pass in garbage
-    )
+        # lats and longs
+        api = TomorrowioV4(api_key, 361.0, 361.0, unit_system="metric", session=session)
+        coordinator = TomorrowioDataUpdateCoordinator(hass, api)
+        hass.data[DOMAIN][api_key] = coordinator
 
-    coordinator = TomorrowioDataUpdateCoordinator(
+    await coordinator.async_setup_entry(entry)
-        hass,
-        entry,
-        api,
-        async_set_update_interval(hass, entry),
-    )
-
-    await coordinator.async_config_entry_first_refresh()
-
-    hass.data[DOMAIN][entry.entry_id] = coordinator
 
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
@@ -227,9 +217,13 @@ async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
         config_entry, PLATFORMS
     )
 
-    hass.data[DOMAIN].pop(config_entry.entry_id)
+    api_key = config_entry.data[CONF_API_KEY]
-    if not hass.data[DOMAIN]:
+    coordinator: TomorrowioDataUpdateCoordinator = hass.data[DOMAIN][api_key]
-        hass.data.pop(DOMAIN)
+    # If this is true, we can remove the coordinator
+    if await coordinator.async_unload_entry(config_entry):
+        hass.data[DOMAIN].pop(api_key)
+        if not hass.data[DOMAIN]:
+            hass.data.pop(DOMAIN)
 
     return unload_ok
 
@@ -237,44 +231,90 @@ async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
 class TomorrowioDataUpdateCoordinator(DataUpdateCoordinator):
     """Define an object to hold Tomorrow.io data."""
 
-    def __init__(
+    def __init__(self, hass: HomeAssistant, api: TomorrowioV4) -> None:
-        self,
-        hass: HomeAssistant,
-        config_entry: ConfigEntry,
-        api: TomorrowioV4,
-        update_interval: timedelta,
-    ) -> None:
         """Initialize."""
-
-        self._config_entry = config_entry
         self._api = api
-        self.name = config_entry.data[CONF_NAME]
         self.data = {CURRENT: {}, FORECASTS: {}}
+        self.entry_id_to_location_dict: dict[str, str] = {}
+        self._coordinator_ready: asyncio.Event | None = None
 
-        super().__init__(
+        super().__init__(hass, _LOGGER, name=f"{DOMAIN}_{self._api.api_key}")
-            hass,
+
-            _LOGGER,
+    def add_entry_to_location_dict(self, entry: ConfigEntry) -> None:
-            name=config_entry.data[CONF_NAME],
+        """Add an entry to the location dict."""
-            update_interval=update_interval,
+        latitude = entry.data[CONF_LOCATION][CONF_LATITUDE]
-        )
+        longitude = entry.data[CONF_LOCATION][CONF_LONGITUDE]
+        self.entry_id_to_location_dict[entry.entry_id] = f"{latitude},{longitude}"
+
+    async def async_setup_entry(self, entry: ConfigEntry) -> None:
+        """Load config entry into coordinator."""
+        # If we haven't loaded any data yet, register all entries with this API key and
+        # get the initial data for all of them. We do this because another config entry
+        # may start setup before we finish setting the initial data and we don't want
+        # to do multiple refreshes on startup.
+        if self._coordinator_ready is None:
+            self._coordinator_ready = asyncio.Event()
+            for entry_ in async_get_entries_by_api_key(self.hass, self._api.api_key):
+                self.add_entry_to_location_dict(entry_)
+            await self.async_config_entry_first_refresh()
+            self._coordinator_ready.set()
+        else:
+            # If we have an event, we need to wait for it to be set before we proceed
+            await self._coordinator_ready.wait()
+            # If we're not getting new data because we already know this entry, we
+            # don't need to schedule a refresh
+            if entry.entry_id in self.entry_id_to_location_dict:
+                return
+            # We need a refresh, but it's going to be a partial refresh so we can
+            # minimize repeat API calls
+            self.add_entry_to_location_dict(entry)
+            await self.async_refresh()
+
+        self.update_interval = async_set_update_interval(self.hass, self._api)
+        self._schedule_refresh()
+
+    async def async_unload_entry(self, entry: ConfigEntry) -> bool | None:
+        """
+        Unload a config entry from coordinator.
+
+        Returns whether coordinator can be removed as well because there are no
+        config entries tied to it anymore.
+        """
+        self.entry_id_to_location_dict.pop(entry.entry_id)
+        self.update_interval = async_set_update_interval(self.hass, self._api, entry)
+        return not self.entry_id_to_location_dict
 
     async def _async_update_data(self) -> dict[str, Any]:
         """Update data via library."""
-        try:
+        data = {}
-            return await self._api.realtime_and_all_forecasts(
+        # If we are refreshing because of a new config entry that's not already in our
-                [
+        # data, we do a partial refresh to avoid wasted API calls.
-                    TMRW_ATTR_TEMPERATURE,
+        if self.data and any(
-                    TMRW_ATTR_HUMIDITY,
+            entry_id not in self.data for entry_id in self.entry_id_to_location_dict
-                    TMRW_ATTR_PRESSURE,
+        ):
-                    TMRW_ATTR_WIND_SPEED,
+            data = self.data
-                    TMRW_ATTR_WIND_DIRECTION,
+
-                    TMRW_ATTR_CONDITION,
+        for entry_id, location in self.entry_id_to_location_dict.items():
-                    TMRW_ATTR_VISIBILITY,
+            if entry_id in data:
-                    TMRW_ATTR_OZONE,
+                continue
-                    TMRW_ATTR_WIND_GUST,
+            entry = self.hass.config_entries.async_get_entry(entry_id)
-                    TMRW_ATTR_CLOUD_COVER,
+            assert entry
-                    TMRW_ATTR_PRECIPITATION_TYPE,
+            try:
-                    *(
+                data[entry_id] = await self._api.realtime_and_all_forecasts(
+                    [
+                        # Weather
+                        TMRW_ATTR_TEMPERATURE,
+                        TMRW_ATTR_HUMIDITY,
+                        TMRW_ATTR_PRESSURE,
+                        TMRW_ATTR_WIND_SPEED,
+                        TMRW_ATTR_WIND_DIRECTION,
+                        TMRW_ATTR_CONDITION,
+                        TMRW_ATTR_VISIBILITY,
+                        TMRW_ATTR_OZONE,
+                        TMRW_ATTR_WIND_GUST,
+                        TMRW_ATTR_CLOUD_COVER,
+                        TMRW_ATTR_PRECIPITATION_TYPE,
+                        # Sensors
                         TMRW_ATTR_CARBON_MONOXIDE,
                         TMRW_ATTR_CHINA_AQI,
                         TMRW_ATTR_CHINA_HEALTH_CONCERN,
@@ -300,26 +340,28 @@ class TomorrowioDataUpdateCoordinator(DataUpdateCoordinator):
                         TMRW_ATTR_SOLAR_GHI,
                         TMRW_ATTR_SULPHUR_DIOXIDE,
                         TMRW_ATTR_WIND_GUST,
-                    ),
+                    ],
-                ],
+                    [
-                [
+                        TMRW_ATTR_TEMPERATURE_LOW,
-                    TMRW_ATTR_TEMPERATURE_LOW,
+                        TMRW_ATTR_TEMPERATURE_HIGH,
-                    TMRW_ATTR_TEMPERATURE_HIGH,
+                        TMRW_ATTR_WIND_SPEED,
-                    TMRW_ATTR_WIND_SPEED,
+                        TMRW_ATTR_WIND_DIRECTION,
-                    TMRW_ATTR_WIND_DIRECTION,
+                        TMRW_ATTR_CONDITION,
-                    TMRW_ATTR_CONDITION,
+                        TMRW_ATTR_PRECIPITATION,
-                    TMRW_ATTR_PRECIPITATION,
+                        TMRW_ATTR_PRECIPITATION_PROBABILITY,
-                    TMRW_ATTR_PRECIPITATION_PROBABILITY,
+                    ],
-                ],
+                    nowcast_timestep=entry.options[CONF_TIMESTEP],
-                nowcast_timestep=self._config_entry.options[CONF_TIMESTEP],
+                    location=location,
-            )
+                )
-        except (
+            except (
-            CantConnectException,
+                CantConnectException,
-            InvalidAPIKeyException,
+                InvalidAPIKeyException,
-            RateLimitedException,
+                RateLimitedException,
-            UnknownException,
+                UnknownException,
-        ) as error:
+            ) as error:
-            raise UpdateFailed from error
+                raise UpdateFailed from error
+
+        return data
 
 
 class TomorrowioEntity(CoordinatorEntity[TomorrowioDataUpdateCoordinator]):
@@ -349,7 +391,8 @@ class TomorrowioEntity(CoordinatorEntity[TomorrowioDataUpdateCoordinator]):
 
         Used for V4 API.
         """
-        return self.coordinator.data.get(CURRENT, {}).get(property_name)
+        entry_id = self._config_entry.entry_id
+        return self.coordinator.data[entry_id].get(CURRENT, {}).get(property_name)
 
     @property
     def attribution(self):

homeassistant/components/tomorrowio/sensor.py

@@ -23,6 +23,7 @@ from homeassistant.const import (
     ATTR_ATTRIBUTION,
     CONCENTRATION_MICROGRAMS_PER_CUBIC_METER,
     CONCENTRATION_PARTS_PER_MILLION,
+    CONF_API_KEY,
     CONF_NAME,
     IRRADIATION_BTUS_PER_HOUR_SQUARE_FOOT,
     IRRADIATION_WATTS_PER_SQUARE_METER,
@@ -286,7 +287,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up a config entry."""
-    coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = hass.data[DOMAIN][config_entry.data[CONF_API_KEY]]
     entities = [
         TomorrowioSensorEntity(hass, config_entry, coordinator, 4, description)
         for description in SENSOR_TYPES

homeassistant/components/tomorrowio/weather.py

@@ -19,6 +19,7 @@ from homeassistant.components.weather import (
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
+    CONF_API_KEY,
     CONF_NAME,
     LENGTH_KILOMETERS,
     LENGTH_MILLIMETERS,
@@ -61,7 +62,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up a config entry."""
-    coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = hass.data[DOMAIN][config_entry.data[CONF_API_KEY]]
 
     entities = [
         TomorrowioWeatherEntity(config_entry, coordinator, 4, forecast_type)
@@ -190,7 +191,11 @@ class TomorrowioWeatherEntity(TomorrowioEntity, WeatherEntity):
     def forecast(self):
         """Return the forecast."""
         # Check if forecasts are available
-        raw_forecasts = self.coordinator.data.get(FORECASTS, {}).get(self.forecast_type)
+        raw_forecasts = (
+            self.coordinator.data.get(self._config_entry.entry_id, {})
+            .get(FORECASTS, {})
+            .get(self.forecast_type)
+        )
         if not raw_forecasts:
             return None
 

tests/components/tomorrowio/conftest.py

@@ -1,6 +1,6 @@
 """Configure py.test."""
 import json
-from unittest.mock import patch
+from unittest.mock import PropertyMock, patch
 
 import pytest
 
@@ -23,8 +23,16 @@ def tomorrowio_config_entry_update_fixture():
     with patch(
         "homeassistant.components.tomorrowio.TomorrowioV4.realtime_and_all_forecasts",
         return_value=json.loads(load_fixture("v4.json", "tomorrowio")),
-    ):
+    ) as mock_update, patch(
-        yield
+        "homeassistant.components.tomorrowio.TomorrowioV4.max_requests_per_day",
+        new_callable=PropertyMock,
+    ) as mock_max_requests_per_day, patch(
+        "homeassistant.components.tomorrowio.TomorrowioV4.num_api_requests",
+        new_callable=PropertyMock,
+    ) as mock_num_api_requests:
+        mock_max_requests_per_day.return_value = 100
+        mock_num_api_requests.return_value = 2
+        yield mock_update
 
 
 @pytest.fixture(name="climacell_config_entry_update")

tests/components/tomorrowio/test_init.py

@@ -1,4 +1,7 @@
 """Tests for Tomorrow.io init."""
+from datetime import timedelta
+from unittest.mock import patch
+
 from homeassistant.components.climacell.const import CONF_TIMESTEP, DOMAIN as CC_DOMAIN
 from homeassistant.components.tomorrowio.config_flow import (
     _get_config_schema,
@@ -17,10 +20,11 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr, entity_registry as er
+from homeassistant.util import dt as dt_util
 
 from .const import MIN_CONFIG
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, async_fire_time_changed
 from tests.components.climacell.const import API_V3_ENTRY_DATA
 
 NEW_NAME = "New Name"
@@ -47,6 +51,69 @@ async def test_load_and_unload(hass: HomeAssistant) -> None:
     assert len(hass.states.async_entity_ids(WEATHER_DOMAIN)) == 0
 
 
+async def test_update_intervals(
+    hass: HomeAssistant, tomorrowio_config_entry_update
+) -> None:
+    """Test coordinator update intervals."""
+    now = dt_util.utcnow()
+    data = _get_config_schema(hass, SOURCE_USER)(MIN_CONFIG)
+    data[CONF_NAME] = "test"
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data=data,
+        options={CONF_TIMESTEP: 1},
+        unique_id=_get_unique_id(hass, data),
+        version=1,
+    )
+    config_entry.add_to_hass(hass)
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+    assert len(tomorrowio_config_entry_update.call_args_list) == 1
+    tomorrowio_config_entry_update.reset_mock()
+
+    # Before the update interval, no updates yet
+    async_fire_time_changed(hass, now + timedelta(minutes=30))
+    await hass.async_block_till_done()
+    assert len(tomorrowio_config_entry_update.call_args_list) == 0
+
+    # On the update interval, we get a new update
+    async_fire_time_changed(hass, now + timedelta(minutes=32))
+    await hass.async_block_till_done()
+    assert len(tomorrowio_config_entry_update.call_args_list) == 1
+    tomorrowio_config_entry_update.reset_mock()
+
+    with patch(
+        "homeassistant.helpers.update_coordinator.utcnow",
+        return_value=now + timedelta(minutes=32),
+    ):
+        # Adding a second config entry should cause the update interval to double
+        config_entry_2 = MockConfigEntry(
+            domain=DOMAIN,
+            data=data,
+            options={CONF_TIMESTEP: 1},
+            unique_id=f"{_get_unique_id(hass, data)}_1",
+            version=1,
+        )
+        config_entry_2.add_to_hass(hass)
+        assert await hass.config_entries.async_setup(config_entry_2.entry_id)
+        await hass.async_block_till_done()
+        assert config_entry.data[CONF_API_KEY] == config_entry_2.data[CONF_API_KEY]
+        # We should get an immediate call once the new config entry is setup for a
+        # partial update
+        assert len(tomorrowio_config_entry_update.call_args_list) == 1
+        tomorrowio_config_entry_update.reset_mock()
+
+    # We should get no new calls on our old interval
+    async_fire_time_changed(hass, now + timedelta(minutes=64))
+    await hass.async_block_till_done()
+    assert len(tomorrowio_config_entry_update.call_args_list) == 0
+
+    # We should get two calls on our new interval, one for each entry
+    async_fire_time_changed(hass, now + timedelta(minutes=96))
+    await hass.async_block_till_done()
+    assert len(tomorrowio_config_entry_update.call_args_list) == 2
+
+
 async def test_climacell_migration_logic(
     hass: HomeAssistant, climacell_config_entry_update
 ) -> None: