Add subentries to WAQI (#148966)

Co-authored-by: Abílio Costa <abmantis@users.noreply.github.com>

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)
+    entry.runtime_data = {}
-    await waqi_coordinator.async_config_entry_first_refresh()
+
-    entry.runtime_data = waqi_coordinator
+    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 @@ from aiowaqi import (
 )
 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 (
+from homeassistant.helpers.selector import LocationSelector
-    LocationSelector,
-    SelectSelector,
-    SelectSelectorConfig,
-)
 
-from .const import CONF_STATION_NUMBER, DOMAIN
+from .const import CONF_STATION_NUMBER, DOMAIN, SUBENTRY_TYPE_STATION
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -54,11 +56,15 @@ 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:
+    @classmethod
-        """Initialize config flow."""
+    @callback
-        self.data: dict[str, Any] = {}
+    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
@@ -66,6 +72,7 @@ class WAQIConfigFlow(ConfigFlow, domain=DOMAIN):
         """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 @@ class WAQIConfigFlow(ConfigFlow, domain=DOMAIN):
                 _LOGGER.exception("Unexpected exception")
                 errors["base"] = "unknown"
             else:
-                self.data = user_input
+                return self.async_create_entry(
-                if user_input[CONF_METHOD] == CONF_MAP:
+                    title="World Air Quality Index",
-                    return await self.async_step_map()
+                    data={
-                return await self.async_step_station_number()
+                        CONF_API_KEY: user_input[CONF_API_KEY],
+                    },
+                )
 
         return self.async_show_form(
             step_id="user",
-            data_schema=vol.Schema(
+            data_schema=vol.Schema({vol.Required(CONF_API_KEY): str}),
-                {
-                    vol.Required(CONF_API_KEY): str,
-                    vol.Required(CONF_METHOD): SelectSelector(
-                        SelectSelectorConfig(
-                            options=[CONF_MAP, CONF_STATION_NUMBER],
-                            translation_key="method",
-                        )
-                    ),
-                }
-            ),
             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 @@ class WAQIConfigFlow(ConfigFlow, domain=DOMAIN):
             data_schema=self.add_suggested_values_to_schema(
                 vol.Schema(
                     {
-                        vol.Required(
+                        vol.Required(CONF_LOCATION): LocationSelector(),
-                            CONF_LOCATION,
-                        ): LocationSelector(),
                     }
                 ),
                 {
@@ -141,12 +151,12 @@ class WAQIConfigFlow(ConfigFlow, domain=DOMAIN):
 
     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 @@ class WAQIConfigFlow(ConfigFlow, domain=DOMAIN):
                 return await self._async_create_entry(measuring_station)
         return self.async_show_form(
             step_id=CONF_STATION_NUMBER,
-            data_schema=vol.Schema(
+            data_schema=vol.Schema({vol.Required(CONF_STATION_NUMBER): int}),
-                {
-                    vol.Required(
-                        CONF_STATION_NUMBER,
-                    ): int,
-                }
-            ),
             errors=errors,
         )
 
     async def _async_create_entry(
         self, measuring_station: WAQIAirQuality
-    ) -> ConfigFlowResult:
+    ) -> SubentryFlowResult:
-        await self.async_set_unique_id(str(measuring_station.station_id))
+        station_id = str(measuring_station.station_id)
-        self._abort_if_unique_id_configured()
+        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 @@ LOGGER = logging.getLogger(__package__)
 
 CONF_STATION_NUMBER = "station_number"
 
-ISSUE_PLACEHOLDER = {"url": "/config/integrations/dashboard/add?domain=waqi"}
+SUBENTRY_TYPE_STATION = "station"

homeassistant/components/waqi/coordinator.py

@@ -6,13 +6,13 @@ from datetime import timedelta
 
 from aiowaqi import WAQIAirQuality, WAQIClient, WAQIError
 
-from homeassistant.config_entries import ConfigEntry
+from homeassistant.config_entries import ConfigEntry, ConfigSubentry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .const import CONF_STATION_NUMBER, DOMAIN, LOGGER
+from .const import CONF_STATION_NUMBER, LOGGER
 
-type WAQIConfigEntry = ConfigEntry[WAQIDataUpdateCoordinator]
+type WAQIConfigEntry = ConfigEntry[dict[str, WAQIDataUpdateCoordinator]]
 
 
 class WAQIDataUpdateCoordinator(DataUpdateCoordinator[WAQIAirQuality]):
@@ -21,22 +21,27 @@ class WAQIDataUpdateCoordinator(DataUpdateCoordinator[WAQIAirQuality]):
     config_entry: WAQIConfigEntry
 
     def __init__(
-        self, hass: HomeAssistant, config_entry: WAQIConfigEntry, client: WAQIClient
+        self,
+        hass: HomeAssistant,
+        config_entry: WAQIConfigEntry,
+        subentry: ConfigSubentry,
+        client: WAQIClient,
     ) -> None:
         """Initialize the WAQI data coordinator."""
         super().__init__(
             hass,
             LOGGER,
             config_entry=config_entry,
-            name=DOMAIN,
+            name=subentry.title,
             update_interval=timedelta(minutes=5),
         )
         self._client = client
+        self.subentry = subentry
 
     async def _async_update_data(self) -> WAQIAirQuality:
         try:
             return await self._client.get_by_station_number(
-                self.config_entry.data[CONF_STATION_NUMBER]
+                self.subentry.data[CONF_STATION_NUMBER]
             )
         except WAQIError as exc:
             raise UpdateFailed from exc

homeassistant/components/waqi/sensor.py

@@ -130,11 +130,14 @@ async def async_setup_entry(
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the WAQI sensor."""
-    coordinator = entry.runtime_data
+    for subentry_id, coordinator in entry.runtime_data.items():
         async_add_entities(
+            (
                 WaqiSensor(coordinator, sensor)
                 for sensor in SENSORS
                 if sensor.available_fn(coordinator.data)
+            ),
+            config_subentry_id=subentry_id,
         )
 
 

homeassistant/components/waqi/strings.json

@@ -3,19 +3,10 @@
     "step": {
       "user": {
         "data": {
-          "api_key": "[%key:common::config_flow::data::api_key%]",
+          "api_key": "[%key:common::config_flow::data::api_key%]"
-          "method": "How do you want to select a measuring station?"
-        }
         },
-      "map": {
+        "data_description": {
-        "description": "Select a location to get the closest measuring station.",
+          "api_key": "API key for the World Air Quality Index"
-        "data": {
-          "location": "[%key:common::config_flow::data::location%]"
-        }
-      },
-      "station_number": {
-        "data": {
-          "station_number": "Measuring station number"
         }
       }
     },
@@ -25,15 +16,44 @@
       "unknown": "[%key:common::config_flow::error::unknown%]"
     },
     "abort": {
-      "already_configured": "[%key:common::config_flow::abort::already_configured_device%]"
+      "already_configured": "[%key:common::config_flow::abort::already_configured_service%]"
     }
   },
-  "selector": {
+  "config_subentries": {
-    "method": {
+    "station": {
-      "options": {
+      "step": {
-        "map": "Select nearest from point on the map",
+        "user": {
-        "station_number": "Enter a station number"
+          "title": "Add measuring station",
+          "description": "How do you want to select a measuring station?",
+          "menu_options": {
+            "map": "[%key:common::config_flow::data::location%]",
+            "station_number": "Measuring station number"
           }
+        },
+        "map": {
+          "data": {
+            "location": "[%key:common::config_flow::data::location%]"
+          },
+          "data_description": {
+            "location": "The location to get the nearest measuring station from"
+          }
+        },
+        "station_number": {
+          "data": {
+            "station_number": "[%key:component::waqi::config_subentries::station::step::user::menu_options::station_number%]"
+          },
+          "data_description": {
+            "station_number": "The number of the measuring station"
+          }
+        }
+      },
+      "abort": {
+        "already_configured": "[%key:common::config_flow::abort::already_configured_service%]"
+      },
+      "initiate_flow": {
+        "user": "Add measuring station"
+      },
+      "entry_type": "Measuring station"
     }
   },
   "entity": {

tests/components/waqi/conftest.py

@@ -7,6 +7,7 @@ from aiowaqi import WAQIAirQuality
 import pytest
 
 from homeassistant.components.waqi.const import CONF_STATION_NUMBER, DOMAIN
+from homeassistant.config_entries import ConfigSubentryData
 from homeassistant.const import CONF_API_KEY
 from homeassistant.core import HomeAssistant
 
@@ -27,9 +28,18 @@ def mock_config_entry() -> MockConfigEntry:
     """Mock config entry."""
     return MockConfigEntry(
         domain=DOMAIN,
-        unique_id="4584",
+        title="WAQI",
+        data={CONF_API_KEY: "asd"},
+        version=2,
+        subentries_data=[
+            ConfigSubentryData(
+                data={CONF_STATION_NUMBER: 4585},
+                subentry_id="ABCDEF",
+                subentry_type="station",
                 title="de Jongweg, Utrecht",
-        data={CONF_API_KEY: "asd", CONF_STATION_NUMBER: 4584},
+                unique_id="4585",
+            )
+        ],
     )
 
 

tests/components/waqi/test_config_flow.py

@@ -1,211 +1,381 @@
 """Test the World Air Quality Index (WAQI) config flow."""
 
-from typing import Any
 from unittest.mock import AsyncMock
 
 from aiowaqi import WAQIAuthenticationError, WAQIConnectionError
 import pytest
 
-from homeassistant.components.waqi.config_flow import CONF_MAP
 from homeassistant.components.waqi.const import CONF_STATION_NUMBER, DOMAIN
-from homeassistant.config_entries import SOURCE_USER
+from homeassistant.config_entries import SOURCE_USER, ConfigSubentryData
 from homeassistant.const import (
+    ATTR_LATITUDE,
+    ATTR_LOCATION,
+    ATTR_LONGITUDE,
     CONF_API_KEY,
-    CONF_LATITUDE,
-    CONF_LOCATION,
-    CONF_LONGITUDE,
-    CONF_METHOD,
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 
+from tests.common import MockConfigEntry
 
-@pytest.mark.parametrize(
+
-    ("method", "payload"),
+@pytest.fixture
-    [
+def second_mock_config_entry() -> MockConfigEntry:
-        (
+    """Mock config entry."""
-            CONF_MAP,
+    return MockConfigEntry(
-            {
+        domain=DOMAIN,
-                CONF_LOCATION: {CONF_LATITUDE: 50.0, CONF_LONGITUDE: 10.0},
+        title="WAQI",
-            },
+        data={CONF_API_KEY: "asdf"},
-        ),
+        version=2,
-        (
+        subentries_data=[
-            CONF_STATION_NUMBER,
+            ConfigSubentryData(
-            {
+                data={CONF_STATION_NUMBER: 4584},
-                CONF_STATION_NUMBER: 4584,
+                subentry_id="ABCDEF",
-            },
+                subentry_type="station",
-        ),
+                title="de Jongweg, Utrecht",
+                unique_id="4584",
+            )
         ],
-)
+    )
-async def test_full_map_flow(
+
-    hass: HomeAssistant,
+
-    mock_setup_entry: AsyncMock,
+async def test_full_flow(
-    mock_waqi: AsyncMock,
+    hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_waqi: AsyncMock
-    method: str,
-    payload: dict[str, Any],
 ) -> None:
-    """Test we get the form."""
+    """Test full flow."""
     result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": SOURCE_USER}
     )
     assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "user"
+    assert not result["errors"]
 
     result = await hass.config_entries.flow.async_configure(
-        result["flow_id"],
+        result["flow_id"], {CONF_API_KEY: "asd"}
-        {CONF_API_KEY: "asd", CONF_METHOD: method},
-    )
-
-    assert result["type"] is FlowResultType.FORM
-    assert result["step_id"] == method
-
-    result = await hass.config_entries.flow.async_configure(
-        result["flow_id"],
-        payload,
     )
 
     assert result["type"] is FlowResultType.CREATE_ENTRY
-    assert result["title"] == "de Jongweg, Utrecht"
+    assert result["title"] == "World Air Quality Index"
-    assert result["data"] == {
+    assert result["data"] == {CONF_API_KEY: "asd"}
-        CONF_API_KEY: "asd",
-        CONF_STATION_NUMBER: 4584,
-    }
-    assert result["result"].unique_id == "4584"
     assert len(mock_setup_entry.mock_calls) == 1
 
 
 @pytest.mark.parametrize(
     ("exception", "error"),
     [
-        (WAQIAuthenticationError(), "invalid_auth"),
+        (WAQIAuthenticationError("Test error"), "invalid_auth"),
-        (WAQIConnectionError(), "cannot_connect"),
+        (WAQIConnectionError("Test error"), "cannot_connect"),
-        (Exception(), "unknown"),
+        (Exception("Test error"), "unknown"),
     ],
 )
-async def test_flow_errors(
+async def test_entry_errors(
     hass: HomeAssistant,
     mock_setup_entry: AsyncMock,
     mock_waqi: AsyncMock,
     exception: Exception,
     error: str,
 ) -> None:
-    """Test we handle errors during configuration."""
+    """Test full flow."""
     result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": SOURCE_USER}
     )
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "user"
+    assert not result["errors"]
 
     mock_waqi.get_by_ip.side_effect = exception
 
     result = await hass.config_entries.flow.async_configure(
-        result["flow_id"],
+        result["flow_id"], {CONF_API_KEY: "asd"}
-        {CONF_API_KEY: "asd", CONF_METHOD: CONF_MAP},
     )
 
     assert result["type"] is FlowResultType.FORM
     assert result["errors"] == {"base": error}
+    assert result["step_id"] == "user"
 
     mock_waqi.get_by_ip.side_effect = None
 
     result = await hass.config_entries.flow.async_configure(
-        result["flow_id"],
+        result["flow_id"], {CONF_API_KEY: "asd"}
-        {CONF_API_KEY: "asd", CONF_METHOD: CONF_MAP},
-    )
-
-    assert result["type"] is FlowResultType.FORM
-    assert result["step_id"] == "map"
-
-    result = await hass.config_entries.flow.async_configure(
-        result["flow_id"],
-        {
-            CONF_LOCATION: {CONF_LATITUDE: 50.0, CONF_LONGITUDE: 10.0},
-        },
     )
 
     assert result["type"] is FlowResultType.CREATE_ENTRY
 
 
-@pytest.mark.parametrize(
+async def test_duplicate_entry(
-    ("method", "payload", "exception", "error"),
-    [
-        (
-            CONF_MAP,
-            {
-                CONF_LOCATION: {CONF_LATITUDE: 50.0, CONF_LONGITUDE: 10.0},
-            },
-            WAQIConnectionError(),
-            "cannot_connect",
-        ),
-        (
-            CONF_MAP,
-            {
-                CONF_LOCATION: {CONF_LATITUDE: 50.0, CONF_LONGITUDE: 10.0},
-            },
-            Exception(),
-            "unknown",
-        ),
-        (
-            CONF_STATION_NUMBER,
-            {
-                CONF_STATION_NUMBER: 4584,
-            },
-            WAQIConnectionError(),
-            "cannot_connect",
-        ),
-        (
-            CONF_STATION_NUMBER,
-            {
-                CONF_STATION_NUMBER: 4584,
-            },
-            Exception(),
-            "unknown",
-        ),
-    ],
-)
-async def test_error_in_second_step(
     hass: HomeAssistant,
     mock_setup_entry: AsyncMock,
     mock_waqi: AsyncMock,
-    method: str,
+    mock_config_entry: MockConfigEntry,
-    payload: dict[str, Any],
-    exception: Exception,
-    error: str,
 ) -> None:
-    """Test we get the form."""
+    """Test duplicate entry handling."""
+    mock_config_entry.add_to_hass(hass)
     result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": SOURCE_USER}
     )
     assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "user"
+    assert not result["errors"]
 
     result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {CONF_API_KEY: "asd"}
+    )
+
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "already_configured"
+
+
+async def test_full_map_flow(
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    mock_waqi: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test we get the form."""
+    mock_config_entry.add_to_hass(hass)
+    result = await hass.config_entries.subentries.async_init(
+        (mock_config_entry.entry_id, "station"),
+        context={"source": SOURCE_USER},
+    )
+    assert result["type"] is FlowResultType.MENU
+    assert result["step_id"] == "user"
+
+    result = await hass.config_entries.subentries.async_configure(
         result["flow_id"],
-        {CONF_API_KEY: "asd", CONF_METHOD: method},
+        {"next_step_id": "map"},
     )
 
     assert result["type"] is FlowResultType.FORM
-    assert result["step_id"] == method
+    assert result["step_id"] == "map"
+    assert not result["errors"]
 
-    mock_waqi.get_by_coordinates.side_effect = exception
+    result = await hass.config_entries.subentries.async_configure(
-    mock_waqi.get_by_station_number.side_effect = exception
-
-    result = await hass.config_entries.flow.async_configure(
         result["flow_id"],
-        payload,
+        {
-    )
+            ATTR_LOCATION: {ATTR_LATITUDE: 50.0, ATTR_LONGITUDE: 10.0},
-
+        },
-    assert result["type"] is FlowResultType.FORM
-    assert result["errors"] == {"base": error}
-
-    mock_waqi.get_by_coordinates.side_effect = None
-    mock_waqi.get_by_station_number.side_effect = None
-
-    result = await hass.config_entries.flow.async_configure(
-        result["flow_id"],
-        payload,
     )
 
     assert result["type"] is FlowResultType.CREATE_ENTRY
     assert result["title"] == "de Jongweg, Utrecht"
-    assert result["data"] == {
+    assert result["data"] == {CONF_STATION_NUMBER: 4584}
-        CONF_API_KEY: "asd",
+    assert list(mock_config_entry.subentries.values())[1].unique_id == "4584"
-        CONF_STATION_NUMBER: 4584,
+
-    }
+
-    assert len(mock_setup_entry.mock_calls) == 1
+@pytest.mark.parametrize(
+    ("exception", "error"),
+    [
+        (WAQIConnectionError("Test error"), "cannot_connect"),
+        (Exception("Test error"), "unknown"),
+    ],
+)
+async def test_map_flow_errors(
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    mock_waqi: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    exception: Exception,
+    error: str,
+) -> None:
+    """Test we get the form."""
+    mock_config_entry.add_to_hass(hass)
+    result = await hass.config_entries.subentries.async_init(
+        (mock_config_entry.entry_id, "station"),
+        context={"source": SOURCE_USER},
+    )
+    assert result["type"] is FlowResultType.MENU
+    assert result["step_id"] == "user"
+
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {"next_step_id": "map"},
+    )
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "map"
+    assert not result["errors"]
+
+    mock_waqi.get_by_coordinates.side_effect = exception
+
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {
+            ATTR_LOCATION: {ATTR_LATITUDE: 50.0, ATTR_LONGITUDE: 10.0},
+        },
+    )
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "map"
+    assert result["errors"] == {"base": error}
+
+    mock_waqi.get_by_coordinates.side_effect = None
+
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {
+            ATTR_LOCATION: {ATTR_LATITUDE: 50.0, ATTR_LONGITUDE: 10.0},
+        },
+    )
+
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+
+
+async def test_map_duplicate(
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    mock_waqi: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    second_mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test duplicate location handling."""
+    mock_config_entry.add_to_hass(hass)
+    second_mock_config_entry.add_to_hass(hass)
+    result = await hass.config_entries.subentries.async_init(
+        (mock_config_entry.entry_id, "station"),
+        context={"source": SOURCE_USER},
+    )
+    assert result["type"] is FlowResultType.MENU
+    assert result["step_id"] == "user"
+
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {"next_step_id": "map"},
+    )
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "map"
+    assert not result["errors"]
+
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {
+            ATTR_LOCATION: {ATTR_LATITUDE: 50.0, ATTR_LONGITUDE: 10.0},
+        },
+    )
+
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "already_configured"
+
+
+async def test_full_station_number_flow(
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    mock_waqi: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test the station number flow."""
+    mock_config_entry.add_to_hass(hass)
+    result = await hass.config_entries.subentries.async_init(
+        (mock_config_entry.entry_id, "station"),
+        context={"source": SOURCE_USER},
+    )
+    assert result["type"] is FlowResultType.MENU
+    assert result["step_id"] == "user"
+
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {"next_step_id": "station_number"},
+    )
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "station_number"
+    assert not result["errors"]
+
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {CONF_STATION_NUMBER: 4584},
+    )
+
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+    assert result["title"] == "de Jongweg, Utrecht"
+    assert result["data"] == {CONF_STATION_NUMBER: 4584}
+    assert list(mock_config_entry.subentries.values())[1].unique_id == "4584"
+
+
+@pytest.mark.parametrize(
+    ("exception", "error"),
+    [
+        (WAQIConnectionError("Test error"), "cannot_connect"),
+        (Exception("Test error"), "unknown"),
+    ],
+)
+async def test_station_number_flow_errors(
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    mock_waqi: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    exception: Exception,
+    error: str,
+) -> None:
+    """Test we get the form."""
+    mock_config_entry.add_to_hass(hass)
+    result = await hass.config_entries.subentries.async_init(
+        (mock_config_entry.entry_id, "station"),
+        context={"source": SOURCE_USER},
+    )
+    assert result["type"] is FlowResultType.MENU
+    assert result["step_id"] == "user"
+
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {"next_step_id": "station_number"},
+    )
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "station_number"
+    assert not result["errors"]
+
+    mock_waqi.get_by_station_number.side_effect = exception
+
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {CONF_STATION_NUMBER: 4584},
+    )
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "station_number"
+    assert result["errors"] == {"base": error}
+
+    mock_waqi.get_by_station_number.side_effect = None
+
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {CONF_STATION_NUMBER: 4584},
+    )
+
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+
+
+async def test_station_number_duplicate(
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    mock_waqi: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    second_mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test duplicate station number handling."""
+    mock_config_entry.add_to_hass(hass)
+    second_mock_config_entry.add_to_hass(hass)
+    result = await hass.config_entries.subentries.async_init(
+        (mock_config_entry.entry_id, "station"),
+        context={"source": SOURCE_USER},
+    )
+    assert result["type"] is FlowResultType.MENU
+    assert result["step_id"] == "user"
+
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {"next_step_id": "station_number"},
+    )
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "station_number"
+    assert not result["errors"]
+
+    result = await hass.config_entries.subentries.async_configure(
+        result["flow_id"],
+        {CONF_STATION_NUMBER: 4584},
+    )
+
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "already_configured"

tests/components/waqi/test_init.py

@@ -1,11 +1,19 @@
 """Test the World Air Quality Index (WAQI) initialization."""
 
-from unittest.mock import AsyncMock
+from typing import Any
+from unittest.mock import AsyncMock, patch
 
 from aiowaqi import WAQIError
+import pytest
 
-from homeassistant.config_entries import ConfigEntryState
+from homeassistant.components.waqi import DOMAIN
+from homeassistant.components.waqi.const import CONF_STATION_NUMBER
+from homeassistant.config_entries import ConfigEntryDisabler, ConfigEntryState
+from homeassistant.const import CONF_API_KEY
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr, entity_registry as er
+from homeassistant.helpers.device_registry import DeviceEntryDisabler
+from homeassistant.helpers.entity_registry import RegistryEntryDisabler
 
 from . import setup_integration
 
@@ -22,3 +30,297 @@ async def test_setup_failed(
     await setup_integration(hass, mock_config_entry)
 
     assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_migration_from_v1(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    mock_setup_entry: AsyncMock,
+) -> None:
+    """Test migration from version 1 to version 2."""
+    # Create a v1 config entry with conversation options and an entity
+    mock_config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={CONF_API_KEY: "1234", CONF_STATION_NUMBER: 4584},
+        version=1,
+        unique_id="4584",
+        title="de Jongweg, Utrecht",
+    )
+    mock_config_entry.add_to_hass(hass)
+    mock_config_entry_2 = MockConfigEntry(
+        domain=DOMAIN,
+        data={CONF_API_KEY: "1234", CONF_STATION_NUMBER: 4585},
+        version=1,
+        unique_id="4585",
+        title="Not de Jongweg, Utrecht",
+    )
+    mock_config_entry_2.add_to_hass(hass)
+
+    device_1 = device_registry.async_get_or_create(
+        config_entry_id=mock_config_entry.entry_id,
+        identifiers={(DOMAIN, "4584")},
+        name="de Jongweg, Utrecht",
+        entry_type=dr.DeviceEntryType.SERVICE,
+    )
+    entity_registry.async_get_or_create(
+        "sensor",
+        DOMAIN,
+        "4584_air_quality",
+        config_entry=mock_config_entry,
+        device_id=device_1.id,
+        suggested_object_id="de_jongweg_utrecht",
+    )
+
+    device_2 = device_registry.async_get_or_create(
+        config_entry_id=mock_config_entry_2.entry_id,
+        identifiers={(DOMAIN, "4585")},
+        name="Not de Jongweg, Utrecht",
+        entry_type=dr.DeviceEntryType.SERVICE,
+    )
+    entity_registry.async_get_or_create(
+        "sensor",
+        DOMAIN,
+        "4585_air_quality",
+        config_entry=mock_config_entry_2,
+        device_id=device_2.id,
+        suggested_object_id="not_de_jongweg_utrecht",
+    )
+
+    await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    entries = hass.config_entries.async_entries(DOMAIN)
+    assert len(entries) == 1
+    entry = entries[0]
+    assert entry.version == 2
+    assert entry.minor_version == 1
+    assert not entry.options
+    assert entry.title == "WAQI"
+    assert len(entry.subentries) == 2
+
+    subentry = list(entry.subentries.values())[0]
+    assert subentry.subentry_type == "station"
+    assert subentry.data[CONF_STATION_NUMBER] == 4584
+    assert subentry.unique_id == "4584"
+    assert subentry.title == "de Jongweg, Utrecht"
+
+    entity = entity_registry.async_get("sensor.de_jongweg_utrecht")
+    assert entity.unique_id == "4584_air_quality"
+    assert entity.config_subentry_id == subentry.subentry_id
+    assert entity.config_entry_id == entry.entry_id
+
+    assert (device := device_registry.async_get_device(identifiers={(DOMAIN, "4584")}))
+    assert device.identifiers == {(DOMAIN, "4584")}
+    assert device.id == device_1.id
+    assert device.config_entries == {mock_config_entry.entry_id}
+    assert device.config_entries_subentries == {
+        mock_config_entry.entry_id: {subentry.subentry_id}
+    }
+
+    subentry = list(entry.subentries.values())[1]
+    assert subentry.subentry_type == "station"
+    assert subentry.data[CONF_STATION_NUMBER] == 4585
+    assert subentry.unique_id == "4585"
+    assert subentry.title == "Not de Jongweg, Utrecht"
+
+    entity = entity_registry.async_get("sensor.not_de_jongweg_utrecht")
+    assert entity.unique_id == "4585_air_quality"
+    assert entity.config_subentry_id == subentry.subentry_id
+    assert entity.config_entry_id == entry.entry_id
+    assert (device := device_registry.async_get_device(identifiers={(DOMAIN, "4585")}))
+    assert device.identifiers == {(DOMAIN, "4585")}
+    assert device.id == device_2.id
+    assert device.config_entries == {mock_config_entry.entry_id}
+    assert device.config_entries_subentries == {
+        mock_config_entry.entry_id: {subentry.subentry_id}
+    }
+
+
+@pytest.mark.parametrize(
+    (
+        "config_entry_disabled_by",
+        "merged_config_entry_disabled_by",
+        "sensor_subentry_data",
+        "main_config_entry",
+    ),
+    [
+        (
+            [ConfigEntryDisabler.USER, None],
+            None,
+            [
+                {
+                    "sensor_entity_id": "sensor.not_de_jongweg_utrecht_air_quality_index",
+                    "device_disabled_by": None,
+                    "entity_disabled_by": None,
+                    "device": 1,
+                },
+                {
+                    "sensor_entity_id": "sensor.de_jongweg_utrecht_air_quality_index",
+                    "device_disabled_by": DeviceEntryDisabler.USER,
+                    "entity_disabled_by": RegistryEntryDisabler.DEVICE,
+                    "device": 0,
+                },
+            ],
+            1,
+        ),
+        (
+            [None, ConfigEntryDisabler.USER],
+            None,
+            [
+                {
+                    "sensor_entity_id": "sensor.de_jongweg_utrecht_air_quality_index",
+                    "device_disabled_by": DeviceEntryDisabler.USER,
+                    "entity_disabled_by": RegistryEntryDisabler.DEVICE,
+                    "device": 0,
+                },
+                {
+                    "sensor_entity_id": "sensor.not_de_jongweg_utrecht_air_quality_index",
+                    "device_disabled_by": None,
+                    "entity_disabled_by": None,
+                    "device": 1,
+                },
+            ],
+            0,
+        ),
+        (
+            [ConfigEntryDisabler.USER, ConfigEntryDisabler.USER],
+            ConfigEntryDisabler.USER,
+            [
+                {
+                    "sensor_entity_id": "sensor.de_jongweg_utrecht_air_quality_index",
+                    "device_disabled_by": DeviceEntryDisabler.CONFIG_ENTRY,
+                    "entity_disabled_by": RegistryEntryDisabler.CONFIG_ENTRY,
+                    "device": 0,
+                },
+                {
+                    "sensor_entity_id": "sensor.not_de_jongweg_utrecht_air_quality_index",
+                    "device_disabled_by": DeviceEntryDisabler.CONFIG_ENTRY,
+                    "entity_disabled_by": None,
+                    "device": 1,
+                },
+            ],
+            0,
+        ),
+    ],
+)
+async def test_migration_from_v1_disabled(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    config_entry_disabled_by: list[ConfigEntryDisabler | None],
+    merged_config_entry_disabled_by: ConfigEntryDisabler | None,
+    sensor_subentry_data: list[dict[str, Any]],
+    main_config_entry: int,
+) -> None:
+    """Test migration where the config entries are disabled."""
+    mock_config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={CONF_API_KEY: "1234", CONF_STATION_NUMBER: 4584},
+        version=1,
+        unique_id="4584",
+        title="de Jongweg, Utrecht",
+        disabled_by=config_entry_disabled_by[0],
+    )
+    mock_config_entry.add_to_hass(hass)
+    mock_config_entry_2 = MockConfigEntry(
+        domain=DOMAIN,
+        data={CONF_API_KEY: "1234", CONF_STATION_NUMBER: 4585},
+        version=1,
+        unique_id="4585",
+        title="Not de Jongweg, Utrecht",
+        disabled_by=config_entry_disabled_by[1],
+    )
+    mock_config_entry_2.add_to_hass(hass)
+    mock_config_entries = [mock_config_entry, mock_config_entry_2]
+
+    device_1 = device_registry.async_get_or_create(
+        config_entry_id=mock_config_entry.entry_id,
+        identifiers={(DOMAIN, mock_config_entry.unique_id)},
+        name=mock_config_entry.title,
+        entry_type=dr.DeviceEntryType.SERVICE,
+        disabled_by=DeviceEntryDisabler.CONFIG_ENTRY,
+    )
+    entity_registry.async_get_or_create(
+        "sensor",
+        DOMAIN,
+        mock_config_entry.unique_id,
+        config_entry=mock_config_entry,
+        device_id=device_1.id,
+        suggested_object_id="de_jongweg_utrecht_air_quality_index",
+        disabled_by=RegistryEntryDisabler.CONFIG_ENTRY,
+    )
+
+    device_2 = device_registry.async_get_or_create(
+        config_entry_id=mock_config_entry_2.entry_id,
+        identifiers={(DOMAIN, mock_config_entry_2.unique_id)},
+        name=mock_config_entry_2.title,
+        entry_type=dr.DeviceEntryType.SERVICE,
+    )
+    entity_registry.async_get_or_create(
+        "sensor",
+        DOMAIN,
+        mock_config_entry_2.unique_id,
+        config_entry=mock_config_entry_2,
+        device_id=device_2.id,
+        suggested_object_id="not_de_jongweg_utrecht_air_quality_index",
+    )
+
+    devices = [device_1, device_2]
+
+    # Run migration
+    with patch(
+        "homeassistant.components.waqi.async_setup_entry",
+        return_value=True,
+    ):
+        await hass.config_entries.async_setup(mock_config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    entries = hass.config_entries.async_entries(DOMAIN)
+    assert len(entries) == 1
+    entry = entries[0]
+    assert entry.disabled_by is merged_config_entry_disabled_by
+    assert entry.version == 2
+    assert entry.minor_version == 1
+    assert not entry.options
+    assert entry.title == "WAQI"
+    assert len(entry.subentries) == 2
+    station_subentries = [
+        subentry
+        for subentry in entry.subentries.values()
+        if subentry.subentry_type == "station"
+    ]
+    assert len(station_subentries) == 2
+    for subentry in station_subentries:
+        assert subentry.data == {CONF_STATION_NUMBER: int(subentry.unique_id)}
+        assert "de Jongweg" in subentry.title
+
+    assert not device_registry.async_get_device(
+        identifiers={(DOMAIN, mock_config_entry.entry_id)}
+    )
+    assert not device_registry.async_get_device(
+        identifiers={(DOMAIN, mock_config_entry_2.entry_id)}
+    )
+
+    for idx, subentry in enumerate(station_subentries):
+        subentry_data = sensor_subentry_data[idx]
+        entity = entity_registry.async_get(subentry_data["sensor_entity_id"])
+        assert entity.unique_id == subentry.unique_id
+        assert entity.config_subentry_id == subentry.subentry_id
+        assert entity.config_entry_id == entry.entry_id
+        assert entity.disabled_by is subentry_data["entity_disabled_by"]
+
+        assert (
+            device := device_registry.async_get_device(
+                identifiers={(DOMAIN, subentry.unique_id)}
+            )
+        )
+        assert device.identifiers == {(DOMAIN, subentry.unique_id)}
+        assert device.id == devices[subentry_data["device"]].id
+        assert device.config_entries == {
+            mock_config_entries[main_config_entry].entry_id
+        }
+        assert device.config_entries_subentries == {
+            mock_config_entries[main_config_entry].entry_id: {subentry.subentry_id}
+        }
+        assert device.disabled_by is subentry_data["device_disabled_by"]