diff --git a/homeassistant/components/tomorrowio/__init__.py b/homeassistant/components/tomorrowio/__init__.py index d8decc1aea3..cef4662d1a4 100644 --- a/homeassistant/components/tomorrowio/__init__.py +++ b/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( - hass: HomeAssistant, current_entry: ConfigEntry -) -> timedelta: - """Recalculate update_interval based on existing Tomorrow.io instances and update them.""" - api_calls = 2 - # 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_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 +def async_get_entries_by_api_key( + hass: HomeAssistant, api_key: str, exclude_entry: ConfigEntry | None = None +) -> list[ConfigEntry]: + """Get all entries for a given API key.""" + return [ + entry for entry in hass.config_entries.async_entries(DOMAIN) - if entry.entry_id != current_entry.entry_id - and entry.data[CONF_API_KEY] == current_entry.data[CONF_API_KEY] + if entry.data[CONF_API_KEY] == api_key + and (exclude_entry is None or exclude_entry != entry) ] - interval = timedelta( - minutes=( - ceil( - (24 * 60 * (len(other_instance_entry_ids) + 1) * api_calls) - / (MAX_REQUESTS_PER_DAY * 0.9) - ) - ) + +@callback +def async_set_update_interval( + hass: HomeAssistant, api: TomorrowioV4, exclude_entry: ConfigEntry | None = None +) -> 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) ) - - 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 + return timedelta(minutes=minutes) @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( - entry.data[CONF_API_KEY], - entry.data[CONF_LOCATION][CONF_LATITUDE], - entry.data[CONF_LOCATION][CONF_LONGITUDE], - unit_system="metric", - session=async_get_clientsession(hass), - ) + api_key = entry.data[CONF_API_KEY] + # If coordinator already exists for this API key, we'll use that, otherwise + # we have to create a new one + if not (coordinator := hass.data[DOMAIN].get(api_key)): + 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( - hass, - entry, - api, - async_set_update_interval(hass, entry), - ) - - await coordinator.async_config_entry_first_refresh() - - hass.data[DOMAIN][entry.entry_id] = coordinator + await coordinator.async_setup_entry(entry) 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) - if not hass.data[DOMAIN]: - hass.data.pop(DOMAIN) + api_key = config_entry.data[CONF_API_KEY] + coordinator: TomorrowioDataUpdateCoordinator = hass.data[DOMAIN][api_key] + # 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__( - self, - hass: HomeAssistant, - config_entry: ConfigEntry, - api: TomorrowioV4, - update_interval: timedelta, - ) -> None: + def __init__(self, hass: HomeAssistant, api: TomorrowioV4) -> 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__( - hass, - _LOGGER, - name=config_entry.data[CONF_NAME], - update_interval=update_interval, - ) + super().__init__(hass, _LOGGER, name=f"{DOMAIN}_{self._api.api_key}") + + def add_entry_to_location_dict(self, entry: ConfigEntry) -> None: + """Add an entry to the location dict.""" + 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: - return await self._api.realtime_and_all_forecasts( - [ - 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, - *( + data = {} + # 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. + if self.data and any( + entry_id not in self.data for entry_id in self.entry_id_to_location_dict + ): + data = self.data + + for entry_id, location in self.entry_id_to_location_dict.items(): + if entry_id in data: + continue + entry = self.hass.config_entries.async_get_entry(entry_id) + assert entry + 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_HIGH, - TMRW_ATTR_WIND_SPEED, - TMRW_ATTR_WIND_DIRECTION, - TMRW_ATTR_CONDITION, - TMRW_ATTR_PRECIPITATION, - TMRW_ATTR_PRECIPITATION_PROBABILITY, - ], - nowcast_timestep=self._config_entry.options[CONF_TIMESTEP], - ) - except ( - CantConnectException, - InvalidAPIKeyException, - RateLimitedException, - UnknownException, - ) as error: - raise UpdateFailed from error + ], + [ + TMRW_ATTR_TEMPERATURE_LOW, + TMRW_ATTR_TEMPERATURE_HIGH, + TMRW_ATTR_WIND_SPEED, + TMRW_ATTR_WIND_DIRECTION, + TMRW_ATTR_CONDITION, + TMRW_ATTR_PRECIPITATION, + TMRW_ATTR_PRECIPITATION_PROBABILITY, + ], + nowcast_timestep=entry.options[CONF_TIMESTEP], + location=location, + ) + except ( + CantConnectException, + InvalidAPIKeyException, + RateLimitedException, + UnknownException, + ) as 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): diff --git a/homeassistant/components/tomorrowio/sensor.py b/homeassistant/components/tomorrowio/sensor.py index d221922df54..ed4ae915c1c 100644 --- a/homeassistant/components/tomorrowio/sensor.py +++ b/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 diff --git a/homeassistant/components/tomorrowio/weather.py b/homeassistant/components/tomorrowio/weather.py index bf687f8bdca..bde6e6b996b 100644 --- a/homeassistant/components/tomorrowio/weather.py +++ b/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 diff --git a/tests/components/tomorrowio/conftest.py b/tests/components/tomorrowio/conftest.py index 65c69209f0e..9c1fa5baa06 100644 --- a/tests/components/tomorrowio/conftest.py +++ b/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")), - ): - yield + ) as mock_update, patch( + "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") diff --git a/tests/components/tomorrowio/test_init.py b/tests/components/tomorrowio/test_init.py index c9914bf95be..27372094092 100644 --- a/tests/components/tomorrowio/test_init.py +++ b/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: