From ae84c7e15d70502c63f4c786486cb153d0b80aef Mon Sep 17 00:00:00 2001 From: Joost Lekkerkerker Date: Thu, 16 Oct 2025 15:11:52 +0200 Subject: [PATCH] Add subentries to WAQI (#148966) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Abílio Costa --- homeassistant/components/waqi/__init__.py | 143 ++++++- homeassistant/components/waqi/config_flow.py | 97 +++-- homeassistant/components/waqi/const.py | 2 +- homeassistant/components/waqi/coordinator.py | 17 +- homeassistant/components/waqi/sensor.py | 15 +- homeassistant/components/waqi/strings.json | 60 ++- tests/components/waqi/conftest.py | 16 +- tests/components/waqi/test_config_flow.py | 414 +++++++++++++------ tests/components/waqi/test_init.py | 306 +++++++++++++- 9 files changed, 862 insertions(+), 208 deletions(-) diff --git a/homeassistant/components/waqi/__init__.py b/homeassistant/components/waqi/__init__.py index 7b1243ed905..ae5ed197b07 100644 --- a/homeassistant/components/waqi/__init__.py +++ b/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, + ) diff --git a/homeassistant/components/waqi/config_flow.py b/homeassistant/components/waqi/config_flow.py index 8ed2dcd8425..d4090e688d9 100644 --- a/homeassistant/components/waqi/config_flow.py +++ b/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 ( - 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,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: - """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 @@ -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 - 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 @@ class WAQIConfigFlow(ConfigFlow, domain=DOMAIN): data_schema=self.add_suggested_values_to_schema( vol.Schema( { - vol.Required( - CONF_LOCATION, - ): LocationSelector(), + vol.Required(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( - { - 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, ) diff --git a/homeassistant/components/waqi/const.py b/homeassistant/components/waqi/const.py index c5ffea20b46..3e48857b85f 100644 --- a/homeassistant/components/waqi/const.py +++ b/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" diff --git a/homeassistant/components/waqi/coordinator.py b/homeassistant/components/waqi/coordinator.py index f40df4a1b89..0c9e624ba66 100644 --- a/homeassistant/components/waqi/coordinator.py +++ b/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 diff --git a/homeassistant/components/waqi/sensor.py b/homeassistant/components/waqi/sensor.py index c887d893c08..cbec9d7476b 100644 --- a/homeassistant/components/waqi/sensor.py +++ b/homeassistant/components/waqi/sensor.py @@ -130,12 +130,15 @@ async def async_setup_entry( async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up the WAQI sensor.""" - coordinator = entry.runtime_data - async_add_entities( - WaqiSensor(coordinator, sensor) - for sensor in SENSORS - if sensor.available_fn(coordinator.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, + ) class WaqiSensor(CoordinatorEntity[WAQIDataUpdateCoordinator], SensorEntity): diff --git a/homeassistant/components/waqi/strings.json b/homeassistant/components/waqi/strings.json index f455e3ead33..96fefe99f58 100644 --- a/homeassistant/components/waqi/strings.json +++ b/homeassistant/components/waqi/strings.json @@ -3,19 +3,10 @@ "step": { "user": { "data": { - "api_key": "[%key:common::config_flow::data::api_key%]", - "method": "How do you want to select a measuring station?" - } - }, - "map": { - "description": "Select a location to get the closest measuring station.", - "data": { - "location": "[%key:common::config_flow::data::location%]" - } - }, - "station_number": { - "data": { - "station_number": "Measuring station number" + "api_key": "[%key:common::config_flow::data::api_key%]" + }, + "data_description": { + "api_key": "API key for the World Air Quality Index" } } }, @@ -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": { - "method": { - "options": { - "map": "Select nearest from point on the map", - "station_number": "Enter a station number" - } + "config_subentries": { + "station": { + "step": { + "user": { + "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": { diff --git a/tests/components/waqi/conftest.py b/tests/components/waqi/conftest.py index bb64fdef097..fa4c3b50e8d 100644 --- a/tests/components/waqi/conftest.py +++ b/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="de Jongweg, Utrecht", - data={CONF_API_KEY: "asd", CONF_STATION_NUMBER: 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", + unique_id="4585", + ) + ], ) diff --git a/tests/components/waqi/test_config_flow.py b/tests/components/waqi/test_config_flow.py index 03759f96ff5..80a992adb97 100644 --- a/tests/components/waqi/test_config_flow.py +++ b/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"), - [ - ( - CONF_MAP, - { - CONF_LOCATION: {CONF_LATITUDE: 50.0, CONF_LONGITUDE: 10.0}, - }, - ), - ( - CONF_STATION_NUMBER, - { - CONF_STATION_NUMBER: 4584, - }, - ), - ], -) -async def test_full_map_flow( - hass: HomeAssistant, - mock_setup_entry: AsyncMock, - mock_waqi: AsyncMock, - method: str, - payload: dict[str, Any], + +@pytest.fixture +def second_mock_config_entry() -> MockConfigEntry: + """Mock config entry.""" + return MockConfigEntry( + domain=DOMAIN, + title="WAQI", + data={CONF_API_KEY: "asdf"}, + version=2, + subentries_data=[ + ConfigSubentryData( + data={CONF_STATION_NUMBER: 4584}, + subentry_id="ABCDEF", + subentry_type="station", + title="de Jongweg, Utrecht", + unique_id="4584", + ) + ], + ) + + +async def test_full_flow( + hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_waqi: AsyncMock ) -> 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"], - {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, + result["flow_id"], {CONF_API_KEY: "asd"} ) assert result["type"] is FlowResultType.CREATE_ENTRY - assert result["title"] == "de Jongweg, Utrecht" - assert result["data"] == { - CONF_API_KEY: "asd", - CONF_STATION_NUMBER: 4584, - } - assert result["result"].unique_id == "4584" + assert result["title"] == "World Air Quality Index" + assert result["data"] == {CONF_API_KEY: "asd"} assert len(mock_setup_entry.mock_calls) == 1 @pytest.mark.parametrize( ("exception", "error"), [ - (WAQIAuthenticationError(), "invalid_auth"), - (WAQIConnectionError(), "cannot_connect"), - (Exception(), "unknown"), + (WAQIAuthenticationError("Test error"), "invalid_auth"), + (WAQIConnectionError("Test error"), "cannot_connect"), + (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"], - {CONF_API_KEY: "asd", CONF_METHOD: CONF_MAP}, + result["flow_id"], {CONF_API_KEY: "asd"} ) 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"], {CONF_API_KEY: "asd"} + ) + + assert result["type"] is FlowResultType.CREATE_ENTRY + + +async def test_duplicate_entry( + hass: HomeAssistant, + mock_setup_entry: AsyncMock, + mock_waqi: AsyncMock, + mock_config_entry: MockConfigEntry, +) -> None: + """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: CONF_MAP}, + {"next_step_id": "map"}, ) assert result["type"] is FlowResultType.FORM assert result["step_id"] == "map" + assert not result["errors"] - result = await hass.config_entries.flow.async_configure( + result = await hass.config_entries.subentries.async_configure( result["flow_id"], { - CONF_LOCATION: {CONF_LATITUDE: 50.0, CONF_LONGITUDE: 10.0}, + ATTR_LOCATION: {ATTR_LATITUDE: 50.0, ATTR_LONGITUDE: 10.0}, + }, + ) + + 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_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 -@pytest.mark.parametrize( - ("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( +async def test_map_duplicate( hass: HomeAssistant, mock_setup_entry: AsyncMock, mock_waqi: AsyncMock, - method: str, - payload: dict[str, Any], - exception: Exception, - error: str, + mock_config_entry: MockConfigEntry, + second_mock_config_entry: MockConfigEntry, ) -> None: - """Test we get the form.""" - result = await hass.config_entries.flow.async_init( - DOMAIN, context={"source": SOURCE_USER} + """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.FORM + assert result["type"] is FlowResultType.MENU + assert result["step_id"] == "user" - result = await hass.config_entries.flow.async_configure( + 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 - mock_waqi.get_by_station_number.side_effect = exception - - result = await hass.config_entries.flow.async_configure( + result = await hass.config_entries.subentries.async_configure( result["flow_id"], - payload, + { + 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["errors"] == {"base": error} + assert result["step_id"] == "station_number" + assert not result["errors"] - 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 = await hass.config_entries.subentries.async_configure( result["flow_id"], - payload, + {CONF_STATION_NUMBER: 4584}, ) assert result["type"] is FlowResultType.CREATE_ENTRY assert result["title"] == "de Jongweg, Utrecht" - assert result["data"] == { - CONF_API_KEY: "asd", - CONF_STATION_NUMBER: 4584, - } - assert len(mock_setup_entry.mock_calls) == 1 + 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" diff --git a/tests/components/waqi/test_init.py b/tests/components/waqi/test_init.py index 7e4487f8ad2..f4e27eee37c 100644 --- a/tests/components/waqi/test_init.py +++ b/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"]