Do not add switch_as_x config entry to source device (#148346)

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 import entity_registry as er
-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 .const import CONF_INVERT, CONF_TARGET_DOMAIN
 
@@ -19,24 +22,14 @@ _LOGGER = logging.getLogger(__name__)
 
 
 @callback
-def async_add_to_device(
+def async_get_parent_device_id(hass: HomeAssistant, entity_id: str) -> str | None:
-    hass: HomeAssistant, entry: ConfigEntry, entity_id: str
+    """Get the parent device id."""
-) -> str | None:
-    """Add our config entry to the tracked entity's device."""
     registry = er.async_get(hass)
-    device_registry = dr.async_get(hass)
-    device_id = None
 
-    if (
+    if not (wrapped_switch := registry.async_get(entity_id)):
-        not (wrapped_switch := registry.async_get(entity_id))
+        return None
-        or not (device_id := wrapped_switch.device_id)
-        or not (device_registry.async_get(device_id))
-    ):
-        return device_id
 
-    device_registry.async_update_device(device_id, add_config_entry_id=entry.entry_id)
+    return wrapped_switch.device_id
-
-    return 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(

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."""

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(
+            self.device_entry = device
-                connections=device.connections,
-                identifiers=device.identifiers,
-            )
         self._attr_entity_category = entity_category
         self._attr_has_entity_name = has_entity_name
         self._attr_name = name

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):
+            device: dev_reg.DeviceEntry | None
-                try:
+            if self.config_entry:
-                    device = dev_reg.async_get(self.hass).async_get_or_create(
+                if device_info := entity.device_info:
-                        config_entry_id=self.config_entry.entry_id,
+                    try:
-                        config_subentry_id=config_subentry_id,
+                        device = dev_reg.async_get(self.hass).async_get_or_create(
-                        **device_info,
+                            config_entry_id=self.config_entry.entry_id,
-                    )
+                            config_subentry_id=config_subentry_id,
-                except dev_reg.DeviceInfoError as exc:
+                            **device_info,
-                    self.logger.error(
+                        )
-                        "%s: Not adding entity with invalid device info: %s",
+                    except dev_reg.DeviceInfoError as exc:
-                        self.platform_name,
+                        self.logger.error(
-                        str(exc),
+                            "%s: Not adding entity with invalid device info: %s",
-                    )
+                            self.platform_name,
-                    entity.add_to_platform_abort()
+                            str(exc),
-                    return
+                        )
+                        entity.add_to_platform_abort()
+                        return
+                else:
+                    device = entity.device_entry
             else:
                 device = None
 

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,
+                source_device_id, remove_config_entry_id=helper_config_entry_id
-                add_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
+        )

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,
+}

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(

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 == []