From 7da2e765d12a84af9e4e2b3a67840564badf5a86 Mon Sep 17 00:00:00 2001 From: Raman Gupta <7243222+raman325@users.noreply.github.com> Date: Fri, 22 Apr 2022 19:10:13 -0400 Subject: [PATCH] Reorganize tomorrowio code to keep it clean (#70454) --- .../components/tomorrowio/__init__.py | 116 ++++++++++-------- 1 file changed, 62 insertions(+), 54 deletions(-) diff --git a/homeassistant/components/tomorrowio/__init__.py b/homeassistant/components/tomorrowio/__init__.py index 38d4117c1c6..d8decc1aea3 100644 --- a/homeassistant/components/tomorrowio/__init__.py +++ b/homeassistant/components/tomorrowio/__init__.py @@ -25,7 +25,7 @@ from homeassistant.const import ( CONF_LONGITUDE, CONF_NAME, ) -from homeassistant.core import HomeAssistant +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.entity import DeviceInfo @@ -84,7 +84,10 @@ _LOGGER = logging.getLogger(__name__) PLATFORMS = [SENSOR_DOMAIN, WEATHER_DOMAIN] -def _set_update_interval(hass: HomeAssistant, current_entry: ConfigEntry) -> timedelta: +@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 @@ -114,6 +117,61 @@ def _set_update_interval(hass: HomeAssistant, current_entry: ConfigEntry) -> tim return interval +@callback +def async_migrate_entry_from_climacell( + hass: HomeAssistant, + dev_reg: dr.DeviceRegistry, + entry: ConfigEntry, + device: dr.DeviceEntry, +) -> None: + """Migrate a config entry from a Climacell entry.""" + # Remove the old config entry ID from the entry data so we don't try this again + # on the next setup + data = entry.data.copy() + old_config_entry_id = data.pop("old_config_entry_id") + hass.config_entries.async_update_entry(entry, data=data) + _LOGGER.debug( + ( + "Setting up imported climacell entry %s for the first time as " + "tomorrowio entry %s" + ), + old_config_entry_id, + entry.entry_id, + ) + + ent_reg = er.async_get(hass) + for entity_entry in er.async_entries_for_config_entry(ent_reg, old_config_entry_id): + old_platform = entity_entry.platform + # In case the API key has changed due to a V3 -> V4 change, we need to + # generate the new entity's unique ID + new_unique_id = ( + f"{entry.data[CONF_API_KEY]}_" + f"{'_'.join(entity_entry.unique_id.split('_')[1:])}" + ) + ent_reg.async_update_entity_platform( + entity_entry.entity_id, + DOMAIN, + new_unique_id=new_unique_id, + new_config_entry_id=entry.entry_id, + new_device_id=device.id, + ) + assert entity_entry + _LOGGER.debug( + "Migrated %s from %s to %s", + entity_entry.entity_id, + old_platform, + DOMAIN, + ) + + # We only have one device in the registry but we will do a loop just in case + for old_device in dr.async_entries_for_config_entry(dev_reg, old_config_entry_id): + if old_device.name_by_user: + dev_reg.async_update_device(device.id, name_by_user=old_device.name_by_user) + + # Remove the old config entry and now the entry is fully migrated + hass.async_create_task(hass.config_entries.async_remove(old_config_entry_id)) + + async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up Tomorrow.io API from a config entry.""" hass.data.setdefault(DOMAIN, {}) @@ -137,57 +195,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: # entities, we need to remove each old entity, creating a new entity in its place # but attached to this entry. if entry.source == SOURCE_IMPORT and "old_config_entry_id" in entry.data: - # Remove the old config entry ID from the entry data so we don't try this again - # on the next setup - data = entry.data.copy() - old_config_entry_id = data.pop("old_config_entry_id") - hass.config_entries.async_update_entry(entry, data=data) - _LOGGER.debug( - ( - "Setting up imported climacell entry %s for the first time as " - "tomorrowio entry %s" - ), - old_config_entry_id, - entry.entry_id, - ) - - ent_reg = er.async_get(hass) - for entity_entry in er.async_entries_for_config_entry( - ent_reg, old_config_entry_id - ): - old_platform = entity_entry.platform - # In case the API key has changed due to a V3 -> V4 change, we need to - # generate the new entity's unique ID - new_unique_id = ( - f"{entry.data[CONF_API_KEY]}_" - f"{'_'.join(entity_entry.unique_id.split('_')[1:])}" - ) - ent_reg.async_update_entity_platform( - entity_entry.entity_id, - DOMAIN, - new_unique_id=new_unique_id, - new_config_entry_id=entry.entry_id, - new_device_id=device.id, - ) - assert entity_entry - _LOGGER.debug( - "Migrated %s from %s to %s", - entity_entry.entity_id, - old_platform, - DOMAIN, - ) - - # We only have one device in the registry but we will do a loop just in case - for old_device in dr.async_entries_for_config_entry( - dev_reg, old_config_entry_id - ): - if old_device.name_by_user: - dev_reg.async_update_device( - device.id, name_by_user=old_device.name_by_user - ) - - # Remove the old config entry and now the entry is fully migrated - hass.async_create_task(hass.config_entries.async_remove(old_config_entry_id)) + async_migrate_entry_from_climacell(hass, dev_reg, entry, device) api = TomorrowioV4( entry.data[CONF_API_KEY], @@ -201,7 +209,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: hass, entry, api, - _set_update_interval(hass, entry), + async_set_update_interval(hass, entry), ) await coordinator.async_config_entry_first_refresh()