From b775ba29553ac59a1fa4006daf2f144c55ca90c1 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Tue, 8 Jul 2025 13:23:28 +0200 Subject: [PATCH] Do not add switch_as_x config entry to source device (#148346) --- .../components/switch_as_x/__init__.py | 42 +++---- .../components/switch_as_x/config_flow.py | 2 +- .../components/switch_as_x/entity.py | 7 +- homeassistant/helpers/entity_platform.py | 34 +++--- homeassistant/helpers/helper_integration.py | 58 +++++++++- tests/components/switch_as_x/__init__.py | 23 ++++ tests/components/switch_as_x/test_init.py | 92 ++++++++++++++-- tests/helpers/test_helper_integration.py | 104 +++++++++++++++++- 8 files changed, 306 insertions(+), 56 deletions(-) diff --git a/homeassistant/components/switch_as_x/__init__.py b/homeassistant/components/switch_as_x/__init__.py index c77eda9b294..b511e2af2b2 100644 --- a/homeassistant/components/switch_as_x/__init__.py +++ b/homeassistant/components/switch_as_x/__init__.py @@ -10,8 +10,11 @@ from homeassistant.components.homeassistant import exposed_entities from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_ENTITY_ID from homeassistant.core import HomeAssistant, callback -from homeassistant.helpers import device_registry as dr, entity_registry as er -from homeassistant.helpers.helper_integration import async_handle_source_entity_changes +from homeassistant.helpers import entity_registry as er +from homeassistant.helpers.helper_integration import ( + async_handle_source_entity_changes, + async_remove_helper_config_entry_from_source_device, +) from .const import CONF_INVERT, CONF_TARGET_DOMAIN @@ -19,24 +22,14 @@ _LOGGER = logging.getLogger(__name__) @callback -def async_add_to_device( - hass: HomeAssistant, entry: ConfigEntry, entity_id: str -) -> str | None: - """Add our config entry to the tracked entity's device.""" +def async_get_parent_device_id(hass: HomeAssistant, entity_id: str) -> str | None: + """Get the parent device id.""" registry = er.async_get(hass) - device_registry = dr.async_get(hass) - device_id = None - if ( - not (wrapped_switch := registry.async_get(entity_id)) - or not (device_id := wrapped_switch.device_id) - or not (device_registry.async_get(device_id)) - ): - return device_id + if not (wrapped_switch := registry.async_get(entity_id)): + return None - device_registry.async_update_device(device_id, add_config_entry_id=entry.entry_id) - - return device_id + return wrapped_switch.device_id async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: @@ -68,9 +61,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: entry.async_on_unload( async_handle_source_entity_changes( hass, + add_helper_config_entry_to_device=False, helper_config_entry_id=entry.entry_id, set_source_entity_id_or_uuid=set_source_entity_id_or_uuid, - source_device_id=async_add_to_device(hass, entry, entity_id), + source_device_id=async_get_parent_device_id(hass, entity_id), source_entity_id_or_uuid=entry.options[CONF_ENTITY_ID], source_entity_removed=source_entity_removed, ) @@ -96,8 +90,18 @@ async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> options = {**config_entry.options} if config_entry.minor_version < 2: options.setdefault(CONF_INVERT, False) + if config_entry.version < 3: + # Remove the switch_as_x config entry from the source device + if source_device_id := async_get_parent_device_id( + hass, options[CONF_ENTITY_ID] + ): + async_remove_helper_config_entry_from_source_device( + hass, + helper_config_entry_id=config_entry.entry_id, + source_device_id=source_device_id, + ) hass.config_entries.async_update_entry( - config_entry, options=options, minor_version=2 + config_entry, options=options, minor_version=3 ) _LOGGER.debug( diff --git a/homeassistant/components/switch_as_x/config_flow.py b/homeassistant/components/switch_as_x/config_flow.py index aa9f1d411ce..cf442256cbe 100644 --- a/homeassistant/components/switch_as_x/config_flow.py +++ b/homeassistant/components/switch_as_x/config_flow.py @@ -58,7 +58,7 @@ class SwitchAsXConfigFlowHandler(SchemaConfigFlowHandler, domain=DOMAIN): options_flow = OPTIONS_FLOW VERSION = 1 - MINOR_VERSION = 2 + MINOR_VERSION = 3 def async_config_entry_title(self, options: Mapping[str, Any]) -> str: """Return config entry title and hide the wrapped entity if registered.""" diff --git a/homeassistant/components/switch_as_x/entity.py b/homeassistant/components/switch_as_x/entity.py index 64bfe712086..7611725d457 100644 --- a/homeassistant/components/switch_as_x/entity.py +++ b/homeassistant/components/switch_as_x/entity.py @@ -15,7 +15,6 @@ from homeassistant.const import ( ) from homeassistant.core import Event, EventStateChangedData, HomeAssistant, callback from homeassistant.helpers import device_registry as dr, entity_registry as er -from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.entity import Entity, ToggleEntity from homeassistant.helpers.event import async_track_state_change_event @@ -48,12 +47,8 @@ class BaseEntity(Entity): if wrapped_switch: name = wrapped_switch.original_name - self._device_id = device_id if device_id and (device := device_registry.async_get(device_id)): - self._attr_device_info = DeviceInfo( - connections=device.connections, - identifiers=device.identifiers, - ) + self.device_entry = device self._attr_entity_category = entity_category self._attr_has_entity_name = has_entity_name self._attr_name = name diff --git a/homeassistant/helpers/entity_platform.py b/homeassistant/helpers/entity_platform.py index 0423a1979bc..e798e85ed02 100644 --- a/homeassistant/helpers/entity_platform.py +++ b/homeassistant/helpers/entity_platform.py @@ -825,21 +825,25 @@ class EntityPlatform: entity.add_to_platform_abort() return - if self.config_entry and (device_info := entity.device_info): - try: - device = dev_reg.async_get(self.hass).async_get_or_create( - config_entry_id=self.config_entry.entry_id, - config_subentry_id=config_subentry_id, - **device_info, - ) - except dev_reg.DeviceInfoError as exc: - self.logger.error( - "%s: Not adding entity with invalid device info: %s", - self.platform_name, - str(exc), - ) - entity.add_to_platform_abort() - return + device: dev_reg.DeviceEntry | None + if self.config_entry: + if device_info := entity.device_info: + try: + device = dev_reg.async_get(self.hass).async_get_or_create( + config_entry_id=self.config_entry.entry_id, + config_subentry_id=config_subentry_id, + **device_info, + ) + except dev_reg.DeviceInfoError as exc: + self.logger.error( + "%s: Not adding entity with invalid device info: %s", + self.platform_name, + str(exc), + ) + entity.add_to_platform_abort() + return + else: + device = entity.device_entry else: device = None diff --git a/homeassistant/helpers/helper_integration.py b/homeassistant/helpers/helper_integration.py index 61bb0bcd45d..d43c1b22a25 100644 --- a/homeassistant/helpers/helper_integration.py +++ b/homeassistant/helpers/helper_integration.py @@ -14,6 +14,7 @@ from .event import async_track_entity_registry_updated_event def async_handle_source_entity_changes( hass: HomeAssistant, *, + add_helper_config_entry_to_device: bool = True, helper_config_entry_id: str, set_source_entity_id_or_uuid: Callable[[str], None], source_device_id: str | None, @@ -88,15 +89,17 @@ def async_handle_source_entity_changes( helper_entity.entity_id, device_id=source_entity_entry.device_id ) - if source_entity_entry.device_id is not None: + if add_helper_config_entry_to_device: + if source_entity_entry.device_id is not None: + device_registry.async_update_device( + source_entity_entry.device_id, + add_config_entry_id=helper_config_entry_id, + ) + device_registry.async_update_device( - source_entity_entry.device_id, - add_config_entry_id=helper_config_entry_id, + source_device_id, remove_config_entry_id=helper_config_entry_id ) - device_registry.async_update_device( - source_device_id, remove_config_entry_id=helper_config_entry_id - ) source_device_id = source_entity_entry.device_id # Reload the config entry so the helper entity is recreated with @@ -111,3 +114,46 @@ def async_handle_source_entity_changes( return async_track_entity_registry_updated_event( hass, source_entity_id, async_registry_updated ) + + +def async_remove_helper_config_entry_from_source_device( + hass: HomeAssistant, + *, + helper_config_entry_id: str, + source_device_id: str, +) -> None: + """Remove helper config entry from source device. + + This is a convenience function to migrate from helpers which added their config + entry to the source device. + """ + device_registry = dr.async_get(hass) + + if ( + not (source_device := device_registry.async_get(source_device_id)) + or helper_config_entry_id not in source_device.config_entries + ): + return + + entity_registry = er.async_get(hass) + helper_entity_entries = er.async_entries_for_config_entry( + entity_registry, helper_config_entry_id + ) + + # Disconnect helper entities from the device to prevent them from + # being removed when the config entry link to the device is removed. + modified_helpers: list[er.RegistryEntry] = [] + for helper in helper_entity_entries: + if helper.device_id != source_device_id: + continue + modified_helpers.append(helper) + entity_registry.async_update_entity(helper.entity_id, device_id=None) + # Remove the helper config entry from the device + device_registry.async_update_device( + source_device_id, remove_config_entry_id=helper_config_entry_id + ) + # Connect the helper entity to the device + for helper in modified_helpers: + entity_registry.async_update_entity( + helper.entity_id, device_id=source_device_id + ) diff --git a/tests/components/switch_as_x/__init__.py b/tests/components/switch_as_x/__init__.py index 2addb832462..dbf1afa54ac 100644 --- a/tests/components/switch_as_x/__init__.py +++ b/tests/components/switch_as_x/__init__.py @@ -1,6 +1,11 @@ """The tests for Switch as X platforms.""" +from homeassistant.components.cover import CoverEntityFeature +from homeassistant.components.fan import FanEntityFeature +from homeassistant.components.light import ATTR_SUPPORTED_COLOR_MODES, ColorMode from homeassistant.components.lock import LockState +from homeassistant.components.siren import SirenEntityFeature +from homeassistant.components.valve import ValveEntityFeature from homeassistant.const import STATE_CLOSED, STATE_OFF, STATE_ON, STATE_OPEN, Platform PLATFORMS_TO_TEST = ( @@ -12,6 +17,15 @@ PLATFORMS_TO_TEST = ( Platform.VALVE, ) +CAPABILITY_MAP = { + Platform.COVER: None, + Platform.FAN: {}, + Platform.LIGHT: {ATTR_SUPPORTED_COLOR_MODES: [ColorMode.ONOFF]}, + Platform.LOCK: None, + Platform.SIREN: None, + Platform.VALVE: None, +} + STATE_MAP = { False: { Platform.COVER: {STATE_ON: STATE_OPEN, STATE_OFF: STATE_CLOSED}, @@ -30,3 +44,12 @@ STATE_MAP = { Platform.VALVE: {STATE_ON: STATE_CLOSED, STATE_OFF: STATE_OPEN}, }, } + +SUPPORTED_FEATURE_MAP = { + Platform.COVER: CoverEntityFeature.OPEN | CoverEntityFeature.CLOSE, + Platform.FAN: FanEntityFeature.TURN_ON | FanEntityFeature.TURN_OFF, + Platform.LIGHT: 0, + Platform.LOCK: 0, + Platform.SIREN: SirenEntityFeature.TURN_ON | SirenEntityFeature.TURN_OFF, + Platform.VALVE: ValveEntityFeature.OPEN | ValveEntityFeature.CLOSE, +} diff --git a/tests/components/switch_as_x/test_init.py b/tests/components/switch_as_x/test_init.py index 2c87b0e3a92..a201cb258d6 100644 --- a/tests/components/switch_as_x/test_init.py +++ b/tests/components/switch_as_x/test_init.py @@ -25,12 +25,12 @@ from homeassistant.const import ( EntityCategory, Platform, ) -from homeassistant.core import Event, HomeAssistant +from homeassistant.core import Event, HomeAssistant, callback from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.helpers.event import async_track_entity_registry_updated_event from homeassistant.setup import async_setup_component -from . import PLATFORMS_TO_TEST +from . import CAPABILITY_MAP, PLATFORMS_TO_TEST, SUPPORTED_FEATURE_MAP from tests.common import MockConfigEntry @@ -79,6 +79,22 @@ def switch_as_x_config_entry( return config_entry +def track_entity_registry_actions( + hass: HomeAssistant, entity_id: str +) -> list[er.EventEntityRegistryUpdatedData]: + """Track entity registry actions for an entity.""" + events = [] + + @callback + def add_event(event: Event[er.EventEntityRegistryUpdatedData]) -> None: + """Add entity registry updated event to the list.""" + events.append(event.data) + + async_track_entity_registry_updated_event(hass, entity_id, add_event) + + return events + + @pytest.mark.parametrize("target_domain", PLATFORMS_TO_TEST) async def test_config_entry_unregistered_uuid( hass: HomeAssistant, target_domain: str @@ -222,7 +238,7 @@ async def test_device_registry_config_entry_1( assert entity_entry.device_id == switch_entity_entry.device_id device_entry = device_registry.async_get(device_entry.id) - assert switch_as_x_config_entry.entry_id in device_entry.config_entries + assert switch_as_x_config_entry.entry_id not in device_entry.config_entries events = [] @@ -304,7 +320,7 @@ async def test_device_registry_config_entry_2( assert entity_entry.device_id == switch_entity_entry.device_id device_entry = device_registry.async_get(device_entry.id) - assert switch_as_x_config_entry.entry_id in device_entry.config_entries + assert switch_as_x_config_entry.entry_id not in device_entry.config_entries events = [] @@ -386,7 +402,7 @@ async def test_device_registry_config_entry_3( assert entity_entry.device_id == switch_entity_entry.device_id device_entry = device_registry.async_get(device_entry.id) - assert switch_as_x_config_entry.entry_id in device_entry.config_entries + assert switch_as_x_config_entry.entry_id not in device_entry.config_entries device_entry_2 = device_registry.async_get(device_entry_2.id) assert switch_as_x_config_entry.entry_id not in device_entry_2.config_entries @@ -413,7 +429,7 @@ async def test_device_registry_config_entry_3( device_entry = device_registry.async_get(device_entry.id) assert switch_as_x_config_entry.entry_id not in device_entry.config_entries device_entry_2 = device_registry.async_get(device_entry_2.id) - assert switch_as_x_config_entry.entry_id in device_entry_2.config_entries + assert switch_as_x_config_entry.entry_id not in device_entry_2.config_entries # Check that the switch_as_x config entry is not removed assert switch_as_x_config_entry.entry_id in hass.config_entries.async_entry_ids() @@ -1083,11 +1099,31 @@ async def test_restore_expose_settings( @pytest.mark.parametrize("target_domain", PLATFORMS_TO_TEST) async def test_migrate( hass: HomeAssistant, + device_registry: dr.DeviceRegistry, entity_registry: er.EntityRegistry, target_domain: Platform, ) -> None: """Test migration.""" - # Setup the config entry + # Switch config entry, device and entity + switch_config_entry = MockConfigEntry() + switch_config_entry.add_to_hass(hass) + + device_entry = device_registry.async_get_or_create( + config_entry_id=switch_config_entry.entry_id, + connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")}, + ) + switch_entity_entry = entity_registry.async_get_or_create( + "switch", + "test", + "unique", + config_entry=switch_config_entry, + device_id=device_entry.id, + original_name="ABC", + suggested_object_id="test", + ) + assert switch_entity_entry.entity_id == "switch.test" + + # Switch_as_x config entry, device and entity config_entry = MockConfigEntry( data={}, domain=DOMAIN, @@ -1100,9 +1136,37 @@ async def test_migrate( minor_version=1, ) config_entry.add_to_hass(hass) + device_registry.async_update_device( + device_entry.id, add_config_entry_id=config_entry.entry_id + ) + switch_as_x_entity_entry = entity_registry.async_get_or_create( + target_domain, + "switch_as_x", + config_entry.entry_id, + capabilities=CAPABILITY_MAP[target_domain], + config_entry=config_entry, + device_id=device_entry.id, + original_name="ABC", + suggested_object_id="abc", + supported_features=SUPPORTED_FEATURE_MAP[target_domain], + ) + entity_registry.async_update_entity_options( + switch_as_x_entity_entry.entity_id, + DOMAIN, + {"entity_id": "switch.test", "invert": False}, + ) + + events = track_entity_registry_actions(hass, switch_as_x_entity_entry.entity_id) + + # Setup the switch_as_x config entry assert await hass.config_entries.async_setup(config_entry.entry_id) await hass.async_block_till_done() + assert set(entity_registry.entities) == { + switch_entity_entry.entity_id, + switch_as_x_entity_entry.entity_id, + } + # Check migration was successful and added invert option assert config_entry.state is ConfigEntryState.LOADED assert config_entry.options == { @@ -1117,6 +1181,20 @@ async def test_migrate( assert hass.states.get(f"{target_domain}.abc") is not None assert entity_registry.async_get(f"{target_domain}.abc") is not None + # Entity removed from device to prevent deletion, then added back to device + assert events == [ + { + "action": "update", + "changes": {"device_id": device_entry.id}, + "entity_id": switch_as_x_entity_entry.entity_id, + }, + { + "action": "update", + "changes": {"device_id": None}, + "entity_id": switch_as_x_entity_entry.entity_id, + }, + ] + @pytest.mark.parametrize("target_domain", PLATFORMS_TO_TEST) async def test_migrate_from_future( diff --git a/tests/helpers/test_helper_integration.py b/tests/helpers/test_helper_integration.py index 47f1b62feb7..91932a51ac2 100644 --- a/tests/helpers/test_helper_integration.py +++ b/tests/helpers/test_helper_integration.py @@ -6,10 +6,13 @@ from unittest.mock import AsyncMock, Mock import pytest from homeassistant.config_entries import ConfigEntry -from homeassistant.core import Event, HomeAssistant +from homeassistant.core import Event, HomeAssistant, callback from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.helpers.event import async_track_entity_registry_updated_event -from homeassistant.helpers.helper_integration import async_handle_source_entity_changes +from homeassistant.helpers.helper_integration import ( + async_handle_source_entity_changes, + async_remove_helper_config_entry_from_source_device, +) from tests.common import ( MockConfigEntry, @@ -184,6 +187,7 @@ def track_entity_registry_actions(hass: HomeAssistant, entity_id: str) -> list[s """Track entity registry actions for an entity.""" events = [] + @callback def add_event(event: Event[er.EventEntityRegistryUpdatedData]) -> None: """Add entity registry updated event to the list.""" events.append(event.data["action"]) @@ -193,6 +197,20 @@ def track_entity_registry_actions(hass: HomeAssistant, entity_id: str) -> list[s return events +def listen_entity_registry_events(hass: HomeAssistant) -> list[str]: + """Track entity registry actions for an entity.""" + events: list[er.EventEntityRegistryUpdatedData] = [] + + @callback + def add_event(event: Event[er.EventEntityRegistryUpdatedData]) -> None: + """Add entity registry updated event to the list.""" + events.append(event.data) + + hass.bus.async_listen(er.EVENT_ENTITY_REGISTRY_UPDATED, add_event) + + return events + + @pytest.mark.parametrize("use_entity_registry_id", [True, False]) @pytest.mark.usefixtures("mock_helper_flow", "mock_helper_integration") async def test_async_handle_source_entity_changes_source_entity_removed( @@ -425,3 +443,85 @@ async def test_async_handle_source_entity_new_entity_id( # Check we got the expected events assert events == [] + + +@pytest.mark.parametrize("use_entity_registry_id", [True, False]) +@pytest.mark.usefixtures("source_entity_entry") +async def test_async_remove_helper_config_entry_from_source_device( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + helper_config_entry: MockConfigEntry, + helper_entity_entry: er.RegistryEntry, + source_device: dr.DeviceEntry, +) -> None: + """Test removing the helper config entry from the source device.""" + # Add the helper config entry to the source device + device_registry.async_update_device( + source_device.id, add_config_entry_id=helper_config_entry.entry_id + ) + + # Create a helper entity entry, not connected to the source device + extra_helper_entity_entry = entity_registry.async_get_or_create( + "sensor", + HELPER_DOMAIN, + f"{helper_config_entry.entry_id}_2", + config_entry=helper_config_entry, + original_name="ABC", + ) + assert extra_helper_entity_entry.entity_id != helper_entity_entry.entity_id + + events = listen_entity_registry_events(hass) + + async_remove_helper_config_entry_from_source_device( + hass, + helper_config_entry_id=helper_config_entry.entry_id, + source_device_id=source_device.id, + ) + + # Check we got the expected events + assert events == [ + { + "action": "update", + "changes": {"device_id": source_device.id}, + "entity_id": helper_entity_entry.entity_id, + }, + { + "action": "update", + "changes": {"device_id": None}, + "entity_id": helper_entity_entry.entity_id, + }, + ] + + +@pytest.mark.parametrize("use_entity_registry_id", [True, False]) +@pytest.mark.usefixtures("source_entity_entry") +async def test_async_remove_helper_config_entry_from_source_device_helper_not_in_device( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, + helper_config_entry: MockConfigEntry, + helper_entity_entry: er.RegistryEntry, + source_device: dr.DeviceEntry, +) -> None: + """Test removing the helper config entry from the source device.""" + # Create a helper entity entry, not connected to the source device + extra_helper_entity_entry = entity_registry.async_get_or_create( + "sensor", + HELPER_DOMAIN, + f"{helper_config_entry.entry_id}_2", + config_entry=helper_config_entry, + original_name="ABC", + ) + assert extra_helper_entity_entry.entity_id != helper_entity_entry.entity_id + + events = listen_entity_registry_events(hass) + + async_remove_helper_config_entry_from_source_device( + hass, + helper_config_entry_id=helper_config_entry.entry_id, + source_device_id=source_device.id, + ) + + # Check we got the expected events + assert events == []