Reorganize tomorrowio code to keep it clean (#70454)

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,29 +117,14 @@ def _set_update_interval(hass: HomeAssistant, current_entry: ConfigEntry) -> tim
     return interval
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+@callback
-    """Set up Tomorrow.io API from a config entry."""
+def async_migrate_entry_from_climacell(
-    hass.data.setdefault(DOMAIN, {})
+    hass: HomeAssistant,
-
+    dev_reg: dr.DeviceRegistry,
-    # Let's precreate the device so that if this is a first time setup for a config
+    entry: ConfigEntry,
-    # entry imported from a ClimaCell entry, we can apply customizations from the old
+    device: dr.DeviceEntry,
-    # device.
+) -> None:
-    dev_reg = dr.async_get(hass)
+    """Migrate a config entry from a Climacell entry."""
-    device = dev_reg.async_get_or_create(
-        config_entry_id=entry.entry_id,
-        identifiers={(DOMAIN, entry.data[CONF_API_KEY])},
-        name=INTEGRATION_NAME,
-        manufacturer=INTEGRATION_NAME,
-        sw_version="v4",
-        entry_type=dr.DeviceEntryType.SERVICE,
-    )
-
-    # If this is an import and we still have the old config entry ID in the entry data,
-    # it means we are setting this entry up for the first time after a migration from
-    # ClimaCell to Tomorrow.io. In order to preserve any customizations on the ClimaCell
-    # 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()
@@ -152,9 +140,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
 
     ent_reg = er.async_get(hass)
-        for entity_entry in er.async_entries_for_config_entry(
+    for entity_entry in er.async_entries_for_config_entry(ent_reg, old_config_entry_id):
-            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
@@ -178,17 +164,39 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         )
 
     # 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(
+    for old_device in dr.async_entries_for_config_entry(dev_reg, old_config_entry_id):
-            dev_reg, old_config_entry_id
-        ):
         if old_device.name_by_user:
-                dev_reg.async_update_device(
+            dev_reg.async_update_device(device.id, name_by_user=old_device.name_by_user)
-                    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, {})
+
+    # Let's precreate the device so that if this is a first time setup for a config
+    # entry imported from a ClimaCell entry, we can apply customizations from the old
+    # device.
+    dev_reg = dr.async_get(hass)
+    device = dev_reg.async_get_or_create(
+        config_entry_id=entry.entry_id,
+        identifiers={(DOMAIN, entry.data[CONF_API_KEY])},
+        name=INTEGRATION_NAME,
+        manufacturer=INTEGRATION_NAME,
+        sw_version="v4",
+        entry_type=dr.DeviceEntryType.SERVICE,
+    )
+
+    # If this is an import and we still have the old config entry ID in the entry data,
+    # it means we are setting this entry up for the first time after a migration from
+    # ClimaCell to Tomorrow.io. In order to preserve any customizations on the ClimaCell
+    # 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:
+        async_migrate_entry_from_climacell(hass, dev_reg, entry, device)
+
     api = TomorrowioV4(
         entry.data[CONF_API_KEY],
         entry.data[CONF_LOCATION][CONF_LATITUDE],
@@ -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()