Update switch_as_x to handle wrapped switch moved to another device (#146387)

* Update switch_as_x to handle wrapped switch moved to another device

* Reload switch_as_x config entry after updating device

* Make sure the switch_as_x entity is not removed

homeassistant/components/switch_as_x/__init__.py

@@ -13,7 +13,7 @@ from homeassistant.core import Event, HomeAssistant, callback, valid_entity_id
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.event import async_track_entity_registry_updated_event
 
-from .const import CONF_INVERT, CONF_TARGET_DOMAIN
+from .const import CONF_INVERT, CONF_TARGET_DOMAIN, DOMAIN
 from .light import LightSwitch
 
 __all__ = ["LightSwitch"]
@@ -81,7 +81,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
                 await hass.config_entries.async_reload(entry.entry_id)
 
         if device_id and "device_id" in data["changes"]:
-            # If the tracked switch is no longer in the device, remove our config entry
+            # Handle the wrapped switch being moved to a different device or removed
             # from the device
             if (
                 not (entity_entry := entity_registry.async_get(data[CONF_ENTITY_ID]))
@@ -91,10 +91,30 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
                 # No need to do any cleanup
                 return
 
+            # The wrapped switch has been moved to a different device, update the
+            # switch_as_x entity and the device entry to include our config entry
+            switch_as_x_entity_id = entity_registry.async_get_entity_id(
+                entry.options[CONF_TARGET_DOMAIN], DOMAIN, entry.entry_id
+            )
+            if switch_as_x_entity_id:
+                # Update the switch_as_x entity to point to the new device (or no device)
+                entity_registry.async_update_entity(
+                    switch_as_x_entity_id, device_id=entity_entry.device_id
+                )
+
+            if entity_entry.device_id is not None:
+                device_registry.async_update_device(
+                    entity_entry.device_id, add_config_entry_id=entry.entry_id
+                )
+
             device_registry.async_update_device(
                 device_id, remove_config_entry_id=entry.entry_id
             )
 
+            # Reload the config entry so the switch_as_x entity is recreated with
+            # correct device info
+            await hass.config_entries.async_reload(entry.entry_id)
+
     entry.async_on_unload(
         async_track_entity_registry_updated_event(
             hass, entity_id, async_registry_updated

tests/components/switch_as_x/test_init.py

@@ -6,6 +6,7 @@ from unittest.mock import patch
 
 import pytest
 
+from homeassistant.components import switch_as_x
 from homeassistant.components.homeassistant import exposed_entities
 from homeassistant.components.lock import LockState
 from homeassistant.components.switch_as_x.config_flow import SwitchAsXConfigFlowHandler
@@ -24,8 +25,9 @@ from homeassistant.const import (
     EntityCategory,
     Platform,
 )
-from homeassistant.core import HomeAssistant
+from homeassistant.core import Event, HomeAssistant
 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
@@ -222,16 +224,39 @@ async def test_device_registry_config_entry_1(
     device_entry = device_registry.async_get(device_entry.id)
     assert switch_as_x_config_entry.entry_id in device_entry.config_entries
 
-    # Remove the wrapped switch's config entry from the device
+    events = []
+
+    def add_event(event: Event[er.EventEntityRegistryUpdatedData]) -> None:
+        """Add entity registry updated event to the list."""
+        events.append(event.data["action"])
+
+    async_track_entity_registry_updated_event(hass, entity_entry.entity_id, add_event)
+
+    # Remove the wrapped switch's config entry from the device, this removes the
+    # wrapped switch
+    with patch(
+        "homeassistant.components.switch_as_x.async_unload_entry",
+        wraps=switch_as_x.async_unload_entry,
+    ) as mock_setup_entry:
         device_registry.async_update_device(
             device_entry.id, remove_config_entry_id=switch_config_entry.entry_id
         )
         await hass.async_block_till_done()
         await hass.async_block_till_done()
+    mock_setup_entry.assert_called_once()
+
     # Check that the switch_as_x config entry is removed from the device
     device_entry = device_registry.async_get(device_entry.id)
     assert switch_as_x_config_entry.entry_id not in device_entry.config_entries
 
+    # Check that the switch_as_x config entry is removed
+    assert (
+        switch_as_x_config_entry.entry_id not in hass.config_entries.async_entry_ids()
+    )
+
+    # Check we got the expected events
+    assert events == ["remove"]
+
 
 @pytest.mark.parametrize("target_domain", PLATFORMS_TO_TEST)
 async def test_device_registry_config_entry_2(
@@ -281,13 +306,121 @@ async def test_device_registry_config_entry_2(
     device_entry = device_registry.async_get(device_entry.id)
     assert switch_as_x_config_entry.entry_id in device_entry.config_entries
 
+    events = []
+
+    def add_event(event: Event[er.EventEntityRegistryUpdatedData]) -> None:
+        """Add entity registry updated event to the list."""
+        events.append(event.data["action"])
+
+    async_track_entity_registry_updated_event(hass, entity_entry.entity_id, add_event)
+
     # Remove the wrapped switch from the device
-    entity_registry.async_update_entity(switch_entity_entry.entity_id, device_id=None)
+    with patch(
+        "homeassistant.components.switch_as_x.async_unload_entry",
+        wraps=switch_as_x.async_unload_entry,
+    ) as mock_setup_entry:
+        entity_registry.async_update_entity(
+            switch_entity_entry.entity_id, device_id=None
+        )
         await hass.async_block_till_done()
+    mock_setup_entry.assert_called_once()
+
     # Check that the switch_as_x config entry is removed from the device
     device_entry = device_registry.async_get(device_entry.id)
     assert switch_as_x_config_entry.entry_id not in device_entry.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()
+
+    # Check we got the expected events
+    assert events == ["update"]
+
+
+@pytest.mark.parametrize("target_domain", PLATFORMS_TO_TEST)
+async def test_device_registry_config_entry_3(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    target_domain: str,
+) -> None:
+    """Test we add our config entry to the tracked switch's device."""
+    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")},
+    )
+    device_entry_2 = device_registry.async_get_or_create(
+        config_entry_id=switch_config_entry.entry_id,
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:FF")},
+    )
+    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",
+    )
+
+    switch_as_x_config_entry = MockConfigEntry(
+        data={},
+        domain=DOMAIN,
+        options={
+            CONF_ENTITY_ID: switch_entity_entry.id,
+            CONF_INVERT: False,
+            CONF_TARGET_DOMAIN: target_domain,
+        },
+        title="ABC",
+        version=SwitchAsXConfigFlowHandler.VERSION,
+        minor_version=SwitchAsXConfigFlowHandler.MINOR_VERSION,
+    )
+
+    switch_as_x_config_entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(switch_as_x_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    entity_entry = entity_registry.async_get(f"{target_domain}.abc")
+    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
+    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
+
+    events = []
+
+    def add_event(event: Event[er.EventEntityRegistryUpdatedData]) -> None:
+        """Add entity registry updated event to the list."""
+        events.append(event.data["action"])
+
+    async_track_entity_registry_updated_event(hass, entity_entry.entity_id, add_event)
+
+    # Move the wrapped switch to another device
+    with patch(
+        "homeassistant.components.switch_as_x.async_unload_entry",
+        wraps=switch_as_x.async_unload_entry,
+    ) as mock_setup_entry:
+        entity_registry.async_update_entity(
+            switch_entity_entry.entity_id, device_id=device_entry_2.id
+        )
+        await hass.async_block_till_done()
+    mock_setup_entry.assert_called_once()
+
+    # Check that the switch_as_x config entry is moved to the other device
+    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
+
+    # 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()
+
+    # Check we got the expected events
+    assert events == ["update"]
+
 
 @pytest.mark.parametrize("target_domain", PLATFORMS_TO_TEST)
 async def test_config_entry_entity_id(