From e43c8b513ef932519923c1832ea4cff14887af32 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Thu, 10 Mar 2022 21:26:35 +0100
Subject: [PATCH] Listen to entity registry events for wrapped switch in
 switch_as_x (#67962)

* Listen to entity registry events for wrapped switch in switch_as_x

* Simplify test
---
 .../components/switch_as_x/__init__.py        | 23 +++++++-
 tests/components/switch_as_x/test_init.py     | 57 +++++++++++++++++++
 2 files changed, 78 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/switch_as_x/__init__.py b/homeassistant/components/switch_as_x/__init__.py
index 6cfea21a349..65b95c59c6d 100644
--- a/homeassistant/components/switch_as_x/__init__.py
+++ b/homeassistant/components/switch_as_x/__init__.py
@@ -7,8 +7,9 @@ import voluptuous as vol
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_ENTITY_ID
-from homeassistant.core import HomeAssistant
+from homeassistant.core import Event, HomeAssistant
 from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers.event import async_track_entity_registry_updated_event
 
 from .light import LightSwitch
 
@@ -23,7 +24,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up a config entry."""
     registry = er.async_get(hass)
     try:
-        er.async_validate_entity_id(registry, entry.options[CONF_ENTITY_ID])
+        entity_id = er.async_validate_entity_id(registry, entry.options[CONF_ENTITY_ID])
     except vol.Invalid:
         # The entity is identified by an unknown entity registry ID
         _LOGGER.error(
@@ -32,6 +33,24 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         )
         return False
 
+    async def async_registry_updated(event: Event) -> None:
+        """Handle entity registry update."""
+        data = event.data
+        if data["action"] == "remove":
+            await hass.config_entries.async_remove(entry.entry_id)
+
+        if data["action"] != "update" or "entity_id" not in data["changes"]:
+            return
+
+        # Entity_id changed, reload the config entry
+        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
+        )
+    )
+
     hass.config_entries.async_setup_platforms(entry, (entry.options["target_domain"],))
     return True
 
diff --git a/tests/components/switch_as_x/test_init.py b/tests/components/switch_as_x/test_init.py
index 5797c66fab5..a8875def0ad 100644
--- a/tests/components/switch_as_x/test_init.py
+++ b/tests/components/switch_as_x/test_init.py
@@ -1,8 +1,11 @@
 """Tests for the Switch as X."""
+from unittest.mock import patch
+
 import pytest
 
 from homeassistant.components.switch_as_x import DOMAIN
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er
 
 from tests.common import MockConfigEntry
 
@@ -25,3 +28,57 @@ async def test_config_entry_unregistered_uuid(hass: HomeAssistant, target_domain
     await hass.async_block_till_done()
 
     assert len(hass.states.async_all()) == 0
+
+
+@pytest.mark.parametrize("target_domain", ("light",))
+async def test_entity_registry_events(hass: HomeAssistant, target_domain):
+    """Test entity registry events are tracked."""
+    registry = er.async_get(hass)
+    registry_entry = registry.async_get_or_create("switch", "test", "unique")
+    switch_entity_id = registry_entry.entity_id
+    hass.states.async_set(switch_entity_id, "on")
+
+    config_entry = MockConfigEntry(
+        data={},
+        domain=DOMAIN,
+        options={"entity_id": registry_entry.id, "target_domain": target_domain},
+        title="ABC",
+    )
+
+    config_entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert hass.states.get(f"{target_domain}.abc").state == "on"
+
+    # Change entity_id
+    new_switch_entity_id = f"{switch_entity_id}_new"
+    registry.async_update_entity(switch_entity_id, new_entity_id=new_switch_entity_id)
+    hass.states.async_set(new_switch_entity_id, "off")
+    await hass.async_block_till_done()
+
+    # Check tracking the new entity_id
+    await hass.async_block_till_done()
+    assert hass.states.get(f"{target_domain}.abc").state == "off"
+
+    # The old entity_id should no longer be tracked
+    hass.states.async_set(switch_entity_id, "on")
+    await hass.async_block_till_done()
+    assert hass.states.get(f"{target_domain}.abc").state == "off"
+
+    # Check changing name does not reload the config entry
+    with patch(
+        "homeassistant.components.switch_as_x.async_unload_entry",
+    ) as mock_setup_entry:
+        registry.async_update_entity(new_switch_entity_id, name="New name")
+        await hass.async_block_till_done()
+    mock_setup_entry.assert_not_called()
+
+    # Check removing the entity removes the config entry
+    registry.async_remove(new_switch_entity_id)
+    await hass.async_block_till_done()
+
+    assert hass.states.get(f"{target_domain}.abc") is None
+    assert registry.async_get(f"{target_domain}.abc") is None
+    assert len(hass.config_entries.async_entries("switch_as_x")) == 0