diff --git a/homeassistant/components/purpleair/__init__.py b/homeassistant/components/purpleair/__init__.py index 7acdbf1fabd..c90f4c9031c 100644 --- a/homeassistant/components/purpleair/__init__.py +++ b/homeassistant/components/purpleair/__init__.py @@ -6,12 +6,10 @@ from aiopurpleair.models.sensors import SensorModel from homeassistant.config_entries import ConfigEntry from homeassistant.const import ATTR_LATITUDE, ATTR_LONGITUDE, Platform from homeassistant.core import HomeAssistant -import homeassistant.helpers.device_registry as dr from homeassistant.helpers.entity import DeviceInfo from homeassistant.helpers.update_coordinator import CoordinatorEntity -from .config_flow import async_remove_sensor_by_device_id -from .const import CONF_LAST_UPDATE_SENSOR_ADD, DOMAIN +from .const import DOMAIN from .coordinator import PurpleAirDataUpdateCoordinator PLATFORMS = [Platform.SENSOR] @@ -32,26 +30,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_handle_entry_update(hass: HomeAssistant, entry: ConfigEntry) -> None: """Handle an options update.""" - if entry.options.get(CONF_LAST_UPDATE_SENSOR_ADD) is True: - # If the last options update was to add a sensor, we reload the config entry: - await hass.config_entries.async_reload(entry.entry_id) - - -async def async_remove_config_entry_device( - hass: HomeAssistant, config_entry: ConfigEntry, device_entry: dr.DeviceEntry -) -> bool: - """Remove a config entry from a device.""" - new_entry_options = async_remove_sensor_by_device_id( - hass, - config_entry, - device_entry.id, - # remove_device is set to False because in this instance, the device has - # already been removed: - remove_device=False, - ) - return hass.config_entries.async_update_entry( - config_entry, options=new_entry_options - ) + await hass.config_entries.async_reload(entry.entry_id) async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: diff --git a/homeassistant/components/purpleair/config_flow.py b/homeassistant/components/purpleair/config_flow.py index 28744e952ab..0b1be019350 100644 --- a/homeassistant/components/purpleair/config_flow.py +++ b/homeassistant/components/purpleair/config_flow.py @@ -1,6 +1,7 @@ """Config flow for PurpleAir integration.""" from __future__ import annotations +import asyncio from collections.abc import Mapping from copy import deepcopy from dataclasses import dataclass, field @@ -14,13 +15,15 @@ import voluptuous as vol from homeassistant import config_entries from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE -from homeassistant.core import HomeAssistant, callback +from homeassistant.core import Event, HomeAssistant, callback from homeassistant.data_entry_flow import FlowResult from homeassistant.helpers import ( aiohttp_client, config_validation as cv, device_registry as dr, + entity_registry as er, ) +from homeassistant.helpers.event import async_track_state_change_event from homeassistant.helpers.selector import ( SelectOptionDict, SelectSelector, @@ -28,7 +31,7 @@ from homeassistant.helpers.selector import ( SelectSelectorMode, ) -from .const import CONF_LAST_UPDATE_SENSOR_ADD, CONF_SENSOR_INDICES, DOMAIN, LOGGER +from .const import CONF_SENSOR_INDICES, DOMAIN, LOGGER CONF_DISTANCE = "distance" CONF_NEARBY_SENSOR_OPTIONS = "nearby_sensor_options" @@ -117,50 +120,6 @@ def async_get_remove_sensor_schema(sensors: list[SelectOptionDict]) -> vol.Schem ) -@callback -def async_get_sensor_index( - hass: HomeAssistant, config_entry: ConfigEntry, device_entry: dr.DeviceEntry -) -> int: - """Get the sensor index related to a config and device entry. - - Note that this method expects that there will always be a single sensor index per - DeviceEntry. - """ - sensor_index = next( - sensor_index - for sensor_index in config_entry.options[CONF_SENSOR_INDICES] - if (DOMAIN, str(sensor_index)) in device_entry.identifiers - ) - - return cast(int, sensor_index) - - -@callback -def async_remove_sensor_by_device_id( - hass: HomeAssistant, - config_entry: ConfigEntry, - device_id: str, - *, - remove_device: bool = True, -) -> dict[str, Any]: - """Remove a sensor and return update config entry options.""" - device_registry = dr.async_get(hass) - device_entry = device_registry.async_get(device_id) - assert device_entry - - removed_sensor_index = async_get_sensor_index(hass, config_entry, device_entry) - options = deepcopy({**config_entry.options}) - options[CONF_LAST_UPDATE_SENSOR_ADD] = False - options[CONF_SENSOR_INDICES].remove(removed_sensor_index) - - if remove_device: - device_registry.async_update_device( - device_entry.id, remove_config_entry_id=config_entry.entry_id - ) - - return options - - @dataclass class ValidationResult: """Define a validation result.""" @@ -407,7 +366,6 @@ class PurpleAirOptionsFlowHandler(config_entries.OptionsFlow): return self.async_abort(reason="already_configured") options = deepcopy({**self.config_entry.options}) - options[CONF_LAST_UPDATE_SENSOR_ADD] = True options[CONF_SENSOR_INDICES].append(sensor_index) return self.async_create_entry(title="", data=options) @@ -432,8 +390,50 @@ class PurpleAirOptionsFlowHandler(config_entries.OptionsFlow): ), ) - new_entry_options = async_remove_sensor_by_device_id( - self.hass, self.config_entry, user_input[CONF_SENSOR_DEVICE_ID] + device_registry = dr.async_get(self.hass) + entity_registry = er.async_get(self.hass) + + device_id = user_input[CONF_SENSOR_DEVICE_ID] + device_entry = cast(dr.DeviceEntry, device_registry.async_get(device_id)) + + # Determine the entity entries that belong to this device. + entity_entries = er.async_entries_for_device( + entity_registry, device_id, include_disabled_entities=True ) - return self.async_create_entry(title="", data=new_entry_options) + device_entities_removed_event = asyncio.Event() + + @callback + def async_device_entity_state_changed(_: Event) -> None: + """Listen and respond when all device entities are removed.""" + if all( + self.hass.states.get(entity_entry.entity_id) is None + for entity_entry in entity_entries + ): + device_entities_removed_event.set() + + # Track state changes for this device's entities and when they're removed, + # finish the flow: + cancel_state_track = async_track_state_change_event( + self.hass, + [entity_entry.entity_id for entity_entry in entity_entries], + async_device_entity_state_changed, + ) + device_registry.async_update_device( + device_id, remove_config_entry_id=self.config_entry.entry_id + ) + await device_entities_removed_event.wait() + + # Once we're done, we can cancel the state change tracker callback: + cancel_state_track() + + # Build new config entry options: + removed_sensor_index = next( + sensor_index + for sensor_index in self.config_entry.options[CONF_SENSOR_INDICES] + if (DOMAIN, str(sensor_index)) in device_entry.identifiers + ) + options = deepcopy({**self.config_entry.options}) + options[CONF_SENSOR_INDICES].remove(removed_sensor_index) + + return self.async_create_entry(title="", data=options) diff --git a/homeassistant/components/purpleair/const.py b/homeassistant/components/purpleair/const.py index 1de915e3545..60f51a9e7dd 100644 --- a/homeassistant/components/purpleair/const.py +++ b/homeassistant/components/purpleair/const.py @@ -5,6 +5,5 @@ DOMAIN = "purpleair" LOGGER = logging.getLogger(__package__) -CONF_LAST_UPDATE_SENSOR_ADD = "last_update_sensor_add" CONF_READ_KEY = "read_key" CONF_SENSOR_INDICES = "sensor_indices" diff --git a/tests/components/purpleair/test_config_flow.py b/tests/components/purpleair/test_config_flow.py index e6a3cea0c20..066706afb50 100644 --- a/tests/components/purpleair/test_config_flow.py +++ b/tests/components/purpleair/test_config_flow.py @@ -215,7 +215,6 @@ async def test_options_add_sensor( ) assert result["type"] == data_entry_flow.FlowResultType.CREATE_ENTRY assert result["data"] == { - "last_update_sensor_add": True, "sensor_indices": [TEST_SENSOR_INDEX1, TEST_SENSOR_INDEX2], } @@ -278,7 +277,6 @@ async def test_options_remove_sensor(hass, config_entry, setup_config_entry): ) assert result["type"] == data_entry_flow.FlowResultType.CREATE_ENTRY assert result["data"] == { - "last_update_sensor_add": False, "sensor_indices": [], }