From fb7da1ba83dd7846726d6febb47fc07092251a57 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 22 Feb 2024 11:39:53 -1000
Subject: [PATCH] Simplify ESPHome entity removal process (#111076)

---
 homeassistant/components/esphome/entity.py    | 12 ++-------
 .../components/esphome/entry_data.py          | 26 ++-----------------
 2 files changed, 4 insertions(+), 34 deletions(-)

diff --git a/homeassistant/components/esphome/entity.py b/homeassistant/components/esphome/entity.py
index 14602077a94..7b06fadb33f 100644
--- a/homeassistant/components/esphome/entity.py
+++ b/homeassistant/components/esphome/entity.py
@@ -65,10 +65,8 @@ def async_static_info_updated(
         device_info = entry_data.device_info
         if TYPE_CHECKING:
             assert device_info is not None
-        hass.async_create_task(
-            entry_data.async_remove_entities(
-                hass, current_infos.values(), device_info.mac_address
-            )
+        entry_data.async_remove_entities(
+            hass, current_infos.values(), device_info.mac_address
         )
 
     # Then update the actual info
@@ -210,12 +208,6 @@ class EsphomeEntity(Entity, Generic[_InfoT, _StateT]):
         key = self._key
         static_info = self._static_info
 
-        self.async_on_remove(
-            entry_data.async_register_key_static_info_remove_callback(
-                static_info,
-                functools.partial(self.async_remove, force_remove=True),
-            )
-        )
         self.async_on_remove(
             async_dispatcher_connect(
                 hass,
diff --git a/homeassistant/components/esphome/entry_data.py b/homeassistant/components/esphome/entry_data.py
index e60de9fb22f..a15f68fd6cc 100644
--- a/homeassistant/components/esphome/entry_data.py
+++ b/homeassistant/components/esphome/entry_data.py
@@ -123,9 +123,6 @@ class RuntimeEntryData:
     entity_info_callbacks: dict[
         type[EntityInfo], list[Callable[[list[EntityInfo]], None]]
     ] = field(default_factory=dict)
-    entity_info_key_remove_callbacks: dict[
-        tuple[type[EntityInfo], int], list[Callable[[], Coroutine[Any, Any, None]]]
-    ] = field(default_factory=dict)
     entity_info_key_updated_callbacks: dict[
         tuple[type[EntityInfo], int], list[Callable[[EntityInfo], None]]
     ] = field(default_factory=dict)
@@ -177,18 +174,6 @@ class RuntimeEntryData:
         """Unsubscribe to when static info is registered."""
         callbacks.remove(callback_)
 
-    @callback
-    def async_register_key_static_info_remove_callback(
-        self,
-        static_info: EntityInfo,
-        callback_: Callable[[], Coroutine[Any, Any, None]],
-    ) -> CALLBACK_TYPE:
-        """Register to receive callbacks when static info is removed for a specific key."""
-        callback_key = (type(static_info), static_info.key)
-        callbacks = self.entity_info_key_remove_callbacks.setdefault(callback_key, [])
-        callbacks.append(callback_)
-        return partial(self._async_unsubscribe_static_key_remove, callbacks, callback_)
-
     @callback
     def _async_unsubscribe_static_key_remove(
         self,
@@ -243,7 +228,8 @@ class RuntimeEntryData:
         """Unsubscribe to assist pipeline updates."""
         self.assist_pipeline_update_callbacks.remove(update_callback)
 
-    async def async_remove_entities(
+    @callback
+    def async_remove_entities(
         self, hass: HomeAssistant, static_infos: Iterable[EntityInfo], mac: str
     ) -> None:
         """Schedule the removal of an entity."""
@@ -255,14 +241,6 @@ class RuntimeEntryData:
             ):
                 ent_reg.async_remove(entry)
 
-        callbacks: list[Coroutine[Any, Any, None]] = []
-        for static_info in static_infos:
-            callback_key = (type(static_info), static_info.key)
-            if key_callbacks := self.entity_info_key_remove_callbacks.get(callback_key):
-                callbacks.extend([callback_() for callback_ in key_callbacks])
-        if callbacks:
-            await asyncio.gather(*callbacks)
-
     @callback
     def async_update_entity_infos(self, static_infos: Iterable[EntityInfo]) -> None:
         """Call static info updated callbacks."""