Add subentries to WAQI (home-assistant#148966)

Co-authored-by: Abílio Costa <[email protected]>

Author: joostlek

homeassistant/components/waqi/__init__.py

@@ -2,32 +2,169 @@
 
 from __future__ import annotations
 
+from types import MappingProxyType
+from typing import TYPE_CHECKING
+
 from aiowaqi import WAQIClient
 
+from homeassistant.config_entries import ConfigEntry, ConfigSubentry
 from homeassistant.const import CONF_API_KEY, Platform
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import (
+    config_validation as cv,
+    device_registry as dr,
+    entity_registry as er,
+)
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.typing import ConfigType
 
+from .const import CONF_STATION_NUMBER, DOMAIN, SUBENTRY_TYPE_STATION
 from .coordinator import WAQIConfigEntry, WAQIDataUpdateCoordinator
 
+CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
 PLATFORMS: list[Platform] = [Platform.SENSOR]
 
 
+async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
+    """Set up WAQI."""
+
+    await async_migrate_integration(hass)
+    return True
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: WAQIConfigEntry) -> bool:
     """Set up World Air Quality Index (WAQI) from a config entry."""
 
     client = WAQIClient(session=async_get_clientsession(hass))
     client.authenticate(entry.data[CONF_API_KEY])
 
-    waqi_coordinator = WAQIDataUpdateCoordinator(hass, entry, client)
-    await waqi_coordinator.async_config_entry_first_refresh()
-    entry.runtime_data = waqi_coordinator
+    entry.runtime_data = {}
+
+    for subentry in entry.subentries.values():
+        if subentry.subentry_type != SUBENTRY_TYPE_STATION:
+            continue
+
+        # Create a coordinator for each station subentry
+        coordinator = WAQIDataUpdateCoordinator(hass, entry, subentry, client)
+        await coordinator.async_config_entry_first_refresh()
+        entry.runtime_data[subentry.subentry_id] = coordinator
+
+    entry.async_on_unload(entry.add_update_listener(async_update_entry))
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True
 
 
+async def async_update_entry(hass: HomeAssistant, entry: WAQIConfigEntry) -> None:
+    """Update entry."""
+    await hass.config_entries.async_reload(entry.entry_id)
+
+
 async def async_unload_entry(hass: HomeAssistant, entry: WAQIConfigEntry) -> bool:
     """Unload a config entry."""
     return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+
+
+async def async_migrate_integration(hass: HomeAssistant) -> None:
+    """Migrate integration entry structure to subentries."""
+
+    # Make sure we get enabled config entries first
+    entries = sorted(
+        hass.config_entries.async_entries(DOMAIN),
+        key=lambda e: e.disabled_by is not None,
+    )
+    if not any(entry.version == 1 for entry in entries):
+        return
+
+    api_keys_entries: dict[str, tuple[ConfigEntry, bool]] = {}
+    entity_registry = er.async_get(hass)
+    device_registry = dr.async_get(hass)
+
+    for entry in entries:
+        subentry = ConfigSubentry(
+            data=MappingProxyType(
+                {CONF_STATION_NUMBER: entry.data[CONF_STATION_NUMBER]}
+            ),
+            subentry_type="station",
+            title=entry.title,
+            unique_id=entry.unique_id,
+        )
+        if entry.data[CONF_API_KEY] not in api_keys_entries:
+            all_disabled = all(
+                e.disabled_by is not None
+                for e in entries
+                if e.data[CONF_API_KEY] == entry.data[CONF_API_KEY]
+            )
+            api_keys_entries[entry.data[CONF_API_KEY]] = (entry, all_disabled)
+
+        parent_entry, all_disabled = api_keys_entries[entry.data[CONF_API_KEY]]
+
+        hass.config_entries.async_add_subentry(parent_entry, subentry)
+
+        entities = er.async_entries_for_config_entry(entity_registry, entry.entry_id)
+        if TYPE_CHECKING:
+            assert entry.unique_id is not None
+        device = device_registry.async_get_device(
+            identifiers={(DOMAIN, entry.unique_id)}
+        )
+
+        for entity_entry in entities:
+            entity_disabled_by = entity_entry.disabled_by
+            if (
+                entity_disabled_by is er.RegistryEntryDisabler.CONFIG_ENTRY
+                and not all_disabled
+            ):
+                # Device and entity registries don't update the disabled_by flag
+                # when moving a device or entity from one config entry to another,
+                # so we need to do it manually.
+                entity_disabled_by = (
+                    er.RegistryEntryDisabler.DEVICE
+                    if device
+                    else er.RegistryEntryDisabler.USER
+                )
+            entity_registry.async_update_entity(
+                entity_entry.entity_id,
+                config_entry_id=parent_entry.entry_id,
+                config_subentry_id=subentry.subentry_id,
+                disabled_by=entity_disabled_by,
+            )
+
+        if device is not None:
+            # Device and entity registries don't update the disabled_by flag when
+            # moving a device or entity from one config entry to another, so we
+            # need to do it manually.
+            device_disabled_by = device.disabled_by
+            if (
+                device.disabled_by is dr.DeviceEntryDisabler.CONFIG_ENTRY
+                and not all_disabled
+            ):
+                device_disabled_by = dr.DeviceEntryDisabler.USER
+            device_registry.async_update_device(
+                device.id,
+                disabled_by=device_disabled_by,
+                add_config_subentry_id=subentry.subentry_id,
+                add_config_entry_id=parent_entry.entry_id,
+            )
+            if parent_entry.entry_id != entry.entry_id:
+                device_registry.async_update_device(
+                    device.id,
+                    remove_config_entry_id=entry.entry_id,
+                )
+            else:
+                device_registry.async_update_device(
+                    device.id,
+                    remove_config_entry_id=entry.entry_id,
+                    remove_config_subentry_id=None,
+                )
+
+        if parent_entry.entry_id != entry.entry_id:
+            await hass.config_entries.async_remove(entry.entry_id)
+        else:
+            hass.config_entries.async_update_entry(
+                entry,
+                title="WAQI",
+                version=2,
+                data={CONF_API_KEY: entry.data[CONF_API_KEY]},
+                unique_id=None,
+            )

homeassistant/components/waqi/config_flow.py

@@ -13,22 +13,24 @@
 )
 import voluptuous as vol
 
-from homeassistant.config_entries import ConfigFlow, ConfigFlowResult
+from homeassistant.config_entries import (
+    ConfigEntry,
+    ConfigFlow,
+    ConfigFlowResult,
+    ConfigSubentryFlow,
+    SubentryFlowResult,
+)
 from homeassistant.const import (
     CONF_API_KEY,
     CONF_LATITUDE,
     CONF_LOCATION,
     CONF_LONGITUDE,
-    CONF_METHOD,
 )
+from homeassistant.core import callback
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.selector import (
-    LocationSelector,
-    SelectSelector,
-    SelectSelectorConfig,
-)
+from homeassistant.helpers.selector import LocationSelector
 
-from .const import CONF_STATION_NUMBER, DOMAIN
+from .const import CONF_STATION_NUMBER, DOMAIN, SUBENTRY_TYPE_STATION
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -54,18 +56,23 @@ async def get_by_station_number(
 class WAQIConfigFlow(ConfigFlow, domain=DOMAIN):
     """Handle a config flow for World Air Quality Index (WAQI)."""
 
-    VERSION = 1
+    VERSION = 2
 
-    def __init__(self) -> None:
-        """Initialize config flow."""
-        self.data: dict[str, Any] = {}
+    @classmethod
+    @callback
+    def async_get_supported_subentry_types(
+        cls, config_entry: ConfigEntry
+    ) -> dict[str, type[ConfigSubentryFlow]]:
+        """Return subentries supported by this handler."""
+        return {SUBENTRY_TYPE_STATION: StationFlowHandler}
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Handle the initial step."""
         errors: dict[str, str] = {}
         if user_input is not None:
+            self._async_abort_entries_match({CONF_API_KEY: user_input[CONF_API_KEY]})
             client = WAQIClient(session=async_get_clientsession(self.hass))
             client.authenticate(user_input[CONF_API_KEY])
             try:
@@ -78,35 +85,40 @@ async def async_step_user(
                 _LOGGER.exception("Unexpected exception")
                 errors["base"] = "unknown"
             else:
-                self.data = user_input
-                if user_input[CONF_METHOD] == CONF_MAP:
-                    return await self.async_step_map()
-                return await self.async_step_station_number()
+                return self.async_create_entry(
+                    title="World Air Quality Index",
+                    data={
+                        CONF_API_KEY: user_input[CONF_API_KEY],
+                    },
+                )
 
         return self.async_show_form(
             step_id="user",
-            data_schema=vol.Schema(
-                {
-                    vol.Required(CONF_API_KEY): str,
-                    vol.Required(CONF_METHOD): SelectSelector(
-                        SelectSelectorConfig(
-                            options=[CONF_MAP, CONF_STATION_NUMBER],
-                            translation_key="method",
-                        )
-                    ),
-                }
-            ),
+            data_schema=vol.Schema({vol.Required(CONF_API_KEY): str}),
             errors=errors,
         )
 
+
+class StationFlowHandler(ConfigSubentryFlow):
+    """Handle subentry flow."""
+
+    async def async_step_user(
+        self, user_input: dict[str, Any] | None = None
+    ) -> SubentryFlowResult:
+        """User flow to create a sensor subentry."""
+        return self.async_show_menu(
+            step_id="user",
+            menu_options=["map", "station_number"],
+        )
+
     async def async_step_map(
         self, user_input: dict[str, Any] | None = None
-    ) -> ConfigFlowResult:
+    ) -> SubentryFlowResult:
         """Add measuring station via map."""
         errors: dict[str, str] = {}
         if user_input is not None:
             client = WAQIClient(session=async_get_clientsession(self.hass))
-            client.authenticate(self.data[CONF_API_KEY])
+            client.authenticate(self._get_entry().data[CONF_API_KEY])
             try:
                 measuring_station = await client.get_by_coordinates(
                     user_input[CONF_LOCATION][CONF_LATITUDE],
@@ -124,9 +136,7 @@ async def async_step_map(
             data_schema=self.add_suggested_values_to_schema(
                 vol.Schema(
                     {
-                        vol.Required(
-                            CONF_LOCATION,
-                        ): LocationSelector(),
+                        vol.Required(CONF_LOCATION): LocationSelector(),
                     }
                 ),
                 {
@@ -141,12 +151,12 @@ async def async_step_map(
 
     async def async_step_station_number(
         self, user_input: dict[str, Any] | None = None
-    ) -> ConfigFlowResult:
+    ) -> SubentryFlowResult:
         """Add measuring station via station number."""
         errors: dict[str, str] = {}
         if user_input is not None:
             client = WAQIClient(session=async_get_clientsession(self.hass))
-            client.authenticate(self.data[CONF_API_KEY])
+            client.authenticate(self._get_entry().data[CONF_API_KEY])
             station_number = user_input[CONF_STATION_NUMBER]
             measuring_station, errors = await get_by_station_number(
                 client, abs(station_number)
@@ -160,25 +170,22 @@ async def async_step_station_number(
                 return await self._async_create_entry(measuring_station)
         return self.async_show_form(
             step_id=CONF_STATION_NUMBER,
-            data_schema=vol.Schema(
-                {
-                    vol.Required(
-                        CONF_STATION_NUMBER,
-                    ): int,
-                }
-            ),
+            data_schema=vol.Schema({vol.Required(CONF_STATION_NUMBER): int}),
             errors=errors,
         )
 
     async def _async_create_entry(
         self, measuring_station: WAQIAirQuality
-    ) -> ConfigFlowResult:
-        await self.async_set_unique_id(str(measuring_station.station_id))
-        self._abort_if_unique_id_configured()
+    ) -> SubentryFlowResult:
+        station_id = str(measuring_station.station_id)
+        for entry in self.hass.config_entries.async_entries(DOMAIN):
+            for subentry in entry.subentries.values():
+                if subentry.unique_id == station_id:
+                    return self.async_abort(reason="already_configured")
         return self.async_create_entry(
             title=measuring_station.city.name,
             data={
-                CONF_API_KEY: self.data[CONF_API_KEY],
                 CONF_STATION_NUMBER: measuring_station.station_id,
             },
+            unique_id=station_id,
         )

homeassistant/components/waqi/const.py

@@ -8,4 +8,4 @@
 
 CONF_STATION_NUMBER = "station_number"
 
-ISSUE_PLACEHOLDER = {"url": "/config/integrations/dashboard/add?domain=waqi"}
+SUBENTRY_TYPE_STATION = "station"