Migrate YouTube channels to config subentries

homeassistant/components/youtube/__init__.py

@@ -1,14 +1,24 @@
 """Support for YouTube."""
 
+from types import MappingProxyType
+
+from homeassistant.config_entries import ConfigEntry, ConfigSubentry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.config_entry_oauth2_flow import (
     OAuth2Session,
     async_get_config_entry_implementation,
 )
 
 from .api import AsyncConfigEntryAuth
+from .const import (
+    ATTR_TITLE,
+    CONF_CHANNEL_ID,
+    CONF_CHANNELS,
+    DOMAIN,
+    SUBENTRY_TYPE_CHANNEL,
+)
 from .coordinator import YouTubeConfigEntry, YouTubeDataUpdateCoordinator
 
 PLATFORMS = [Platform.SENSOR]
@@ -20,13 +30,18 @@ async def async_setup_entry(hass: HomeAssistant, entry: YouTubeConfigEntry) -> b
     session = OAuth2Session(hass, entry, implementation)
     auth = AsyncConfigEntryAuth(hass, session)
     await auth.check_and_refresh_token()
-    coordinator = YouTubeDataUpdateCoordinator(hass, entry, auth)
 
+    coordinator = YouTubeDataUpdateCoordinator(hass, entry, auth)
     await coordinator.async_config_entry_first_refresh()
 
-    await delete_devices(hass, entry, coordinator)
+    data = coordinator.data
+    for subentry in entry.get_subentries_of_type(SUBENTRY_TYPE_CHANNEL):
+        channel = data.get(subentry.data[CONF_CHANNEL_ID])
+        if channel is not None and (title := channel[ATTR_TITLE]) != subentry.title:
+            hass.config_entries.async_update_subentry(entry, subentry, title=title)
 
     entry.runtime_data = coordinator
+    entry.async_on_unload(entry.add_update_listener(async_update_listener))
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True
@@ -37,15 +52,90 @@ async def async_unload_entry(hass: HomeAssistant, entry: YouTubeConfigEntry) ->
     return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 
-async def delete_devices(
-    hass: HomeAssistant,
-    entry: YouTubeConfigEntry,
-    coordinator: YouTubeDataUpdateCoordinator,
-) -> None:
-    """Delete all devices created by integration."""
-    channel_ids = list(coordinator.data)
+async def async_update_listener(hass: HomeAssistant, entry: YouTubeConfigEntry) -> None:
+    """Reload the config entry when channels are added or removed."""
+    subentry_ids = {
+        subentry.subentry_id
+        for subentry in entry.get_subentries_of_type(SUBENTRY_TYPE_CHANNEL)
+    }
+    if subentry_ids == entry.runtime_data.subentry_ids:
+        # Token refreshes update the entry data only; the coordinator
+        # uses the refreshed token without needing a reload.
+        return
+    await hass.config_entries.async_reload(entry.entry_id)
+
+
+async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Migrate old entries to the subentry structure."""
     device_registry = dr.async_get(hass)
-    dev_entries = dr.async_entries_for_config_entry(device_registry, entry.entry_id)
-    for dev_entry in dev_entries:
-        if any(identifier[1] in channel_ids for identifier in dev_entry.identifiers):
-            device_registry.async_remove_device(dev_entry.id)
+    entity_registry = er.async_get(hass)
+    prefix = f"{entry.entry_id}_"
+    channel_ids = dict.fromkeys(entry.options.get(CONF_CHANNELS, []))
+
+    subentries: dict[str, ConfigSubentry] = {}
+    for channel_id in channel_ids:
+        device = device_registry.async_get_device_by_identifier(
+            (DOMAIN, f"{prefix}{channel_id}"), entry.entry_id
+        )
+        title = channel_id
+        if device is not None and device.name is not None:
+            # Prefer the channel name of the existing device so the title
+            # survives even if the channel can no longer be fetched.
+            title = device.name
+        subentry = ConfigSubentry(
+            data=MappingProxyType({CONF_CHANNEL_ID: channel_id}),
+            subentry_type=SUBENTRY_TYPE_CHANNEL,
+            title=title,
+            unique_id=channel_id,
+        )
+        hass.config_entries.async_add_subentry(entry, subentry)
+        subentries[channel_id] = subentry
+
+    # Attach the entities of tracked channels to their subentry and remove
+    # entities left behind by channels which are no longer tracked.
+    channel_prefixes = {
+        f"{prefix}{channel_id}_": channel_id for channel_id in channel_ids
+    }
+    for entity_entry in er.async_entries_for_config_entry(
+        entity_registry, entry.entry_id
+    ):
+        channel_subentry = next(
+            (
+                subentries[channel_id]
+                for channel_prefix, channel_id in channel_prefixes.items()
+                if entity_entry.unique_id.startswith(channel_prefix)
+            ),
+            None,
+        )
+        if channel_subentry is None:
+            entity_registry.async_remove(entity_entry.entity_id)
+        else:
+            entity_registry.async_update_entity(
+                entity_entry.entity_id,
+                config_subentry_id=channel_subentry.subentry_id,
+            )
+
+    # Move the devices of tracked channels to their subentry and remove
+    # devices left behind by channels which are no longer tracked.
+    for device_entry in dr.async_entries_for_config_entry(
+        device_registry, entry.entry_id
+    ):
+        channel_id = next(
+            (
+                identifier[1].removeprefix(prefix)
+                for identifier in device_entry.identifiers
+                if identifier[0] == DOMAIN and identifier[1].startswith(prefix)
+            ),
+            None,
+        )
+        if channel_id is not None and channel_id in subentries:
+            device_registry.async_update_device(
+                device_entry.id,
+                new_identifiers={(DOMAIN, channel_id)},
+                new_config_subentry_id=subentries[channel_id].subentry_id,
+            )
+        else:
+            device_registry.async_remove_device(device_entry.id)
+
+    hass.config_entries.async_update_entry(entry, version=2, options={})
+    return True