Fix esphome not removing entities when static info changes (home-assi…

…stant#95202)

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>

homeassistant/components/esphome/__init__.py

@@ -675,6 +675,7 @@ async def _cleanup_instance(
     data.disconnect_callbacks = []
     for cleanup_callback in data.cleanup_callbacks:
         cleanup_callback()
+    await data.async_cleanup()
     await data.client.disconnect()
     return data
 

homeassistant/components/esphome/alarm_control_panel.py

@@ -73,7 +73,6 @@ async def async_setup_entry(
         hass,
         entry,
         async_add_entities,
-        component_key="alarm_control_panel",
         info_type=AlarmControlPanelInfo,
         entity_type=EsphomeAlarmControlPanel,
         state_type=AlarmControlPanelEntityState,

homeassistant/components/esphome/binary_sensor.py

@@ -29,7 +29,6 @@ async def async_setup_entry(
         hass,
         entry,
         async_add_entities,
-        component_key="binary_sensor",
         info_type=BinarySensorInfo,
         entity_type=EsphomeBinarySensor,
         state_type=BinarySensorState,

homeassistant/components/esphome/button.py

@@ -23,7 +23,6 @@ async def async_setup_entry(
         hass,
         entry,
         async_add_entities,
-        component_key="button",
         info_type=ButtonInfo,
         entity_type=EsphomeButton,
         state_type=EntityState,

homeassistant/components/esphome/camera.py

@@ -27,7 +27,6 @@ async def async_setup_entry(
         hass,
         entry,
         async_add_entities,
-        component_key="camera",
         info_type=CameraInfo,
         entity_type=EsphomeCamera,
         state_type=CameraState,

homeassistant/components/esphome/climate.py

@@ -72,7 +72,6 @@ async def async_setup_entry(
         hass,
         entry,
         async_add_entities,
-        component_key="climate",
         info_type=ClimateInfo,
         entity_type=EsphomeClimateEntity,
         state_type=ClimateState,

homeassistant/components/esphome/cover.py

@@ -32,7 +32,6 @@ async def async_setup_entry(
         hass,
         entry,
         async_add_entities,
-        component_key="cover",
         info_type=CoverInfo,
         entity_type=EsphomeCover,
         state_type=CoverState,

homeassistant/components/esphome/diagnostics.py

@@ -1,7 +1,7 @@
 """Diagnostics support for ESPHome."""
 from __future__ import annotations
 
-from typing import Any, cast
+from typing import Any
 
 from homeassistant.components.bluetooth import async_scanner_by_source
 from homeassistant.components.diagnostics import async_redact_data
@@ -28,7 +28,6 @@ async def async_get_config_entry_diagnostics(
     entry_data = DomainData.get(hass).get_entry_data(config_entry)
 
     if (storage_data := await entry_data.store.async_load()) is not None:
-        storage_data = cast("dict[str, Any]", storage_data)
         diag["storage_data"] = storage_data
 
     if config_entry.unique_id and (

homeassistant/components/esphome/domain_data.py

@@ -12,10 +12,9 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.json import JSONEncoder
-from homeassistant.helpers.storage import Store
 
 from .const import DOMAIN
-from .entry_data import RuntimeEntryData
+from .entry_data import ESPHomeStorage, RuntimeEntryData
 
 STORAGE_VERSION = 1
 MAX_CACHED_SERVICES = 128
@@ -26,7 +25,7 @@ class DomainData:
     """Define a class that stores global esphome data in hass.data[DOMAIN]."""
 
     _entry_datas: dict[str, RuntimeEntryData] = field(default_factory=dict)
-    _stores: dict[str, Store] = field(default_factory=dict)
+    _stores: dict[str, ESPHomeStorage] = field(default_factory=dict)
     _gatt_services_cache: MutableMapping[int, BleakGATTServiceCollection] = field(
         default_factory=lambda: LRU(MAX_CACHED_SERVICES)
     )
@@ -83,11 +82,13 @@ def is_entry_loaded(self, entry: ConfigEntry) -> bool:
         """Check whether the given entry is loaded."""
         return entry.entry_id in self._entry_datas
 
-    def get_or_create_store(self, hass: HomeAssistant, entry: ConfigEntry) -> Store:
+    def get_or_create_store(
+        self, hass: HomeAssistant, entry: ConfigEntry
+    ) -> ESPHomeStorage:
         """Get or create a Store instance for the given config entry."""
         return self._stores.setdefault(
             entry.entry_id,
-            Store(
+            ESPHomeStorage(
                 hass, STORAGE_VERSION, f"esphome.{entry.entry_id}", encoder=JSONEncoder
             ),
         )

homeassistant/components/esphome/entity.py

@@ -27,7 +27,6 @@
 import homeassistant.helpers.device_registry as dr
 from homeassistant.helpers.dispatcher import (
     async_dispatcher_connect,
-    async_dispatcher_send,
 )
 from homeassistant.helpers.entity import DeviceInfo, Entity
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -49,7 +48,6 @@ async def platform_async_setup_entry(
     entry: ConfigEntry,
     async_add_entities: AddEntitiesCallback,
     *,
-    component_key: str,
     info_type: type[_InfoT],
     entity_type: type[_EntityT],
     state_type: type[_StateT],
@@ -60,41 +58,34 @@ async def platform_async_setup_entry(
     info and state updates.
     """
     entry_data: RuntimeEntryData = DomainData.get(hass).get_entry_data(entry)
-    entry_data.info[component_key] = {}
-    entry_data.old_info[component_key] = {}
+    entry_data.info[info_type] = {}
     entry_data.state.setdefault(state_type, {})
 
     @callback
     def async_list_entities(infos: list[EntityInfo]) -> None:
         """Update entities of this platform when entities are listed."""
-        old_infos = entry_data.info[component_key]
+        current_infos = entry_data.info[info_type]
         new_infos: dict[int, EntityInfo] = {}
         add_entities: list[_EntityT] = []
+
         for info in infos:
-            if info.key in old_infos:
-                # Update existing entity
-                old_infos.pop(info.key)
-            else:
+            if not current_infos.pop(info.key, None):
                 # Create new entity
-                entity = entity_type(entry_data, component_key, info, state_type)
+                entity = entity_type(entry_data, info, state_type)
                 add_entities.append(entity)
             new_infos[info.key] = info
 
-        # Remove old entities
-        for info in old_infos.values():
-            entry_data.async_remove_entity(hass, component_key, info.key)
+        # Anything still in current_infos is now gone
+        if current_infos:
+            hass.async_create_task(
+                entry_data.async_remove_entities(current_infos.values())
+            )
 
-        # First copy the now-old info into the backup object
-        entry_data.old_info[component_key] = entry_data.info[component_key]
         # Then update the actual info
-        entry_data.info[component_key] = new_infos
+        entry_data.info[info_type] = new_infos
 
-        for key, new_info in new_infos.items():
-            async_dispatcher_send(
-                hass,
-                entry_data.signal_component_key_static_info_updated(component_key, key),
-                new_info,
-            )
+        if new_infos:
+            entry_data.async_update_entity_infos(new_infos.values())
 
         if add_entities:
             # Add entities to Home Assistant
@@ -154,14 +145,12 @@ class EsphomeEntity(Entity, Generic[_InfoT, _StateT]):
     def __init__(
         self,
         entry_data: RuntimeEntryData,
-        component_key: str,
         entity_info: EntityInfo,
         state_type: type[_StateT],
     ) -> None:
         """Initialize."""
         self._entry_data = entry_data
         self._on_entry_data_changed()
-        self._component_key = component_key
         self._key = entity_info.key
         self._state_type = state_type
         self._on_static_info_update(entity_info)
@@ -178,13 +167,11 @@ async def async_added_to_hass(self) -> None:
         """Register callbacks."""
         entry_data = self._entry_data
         hass = self.hass
-        component_key = self._component_key
         key = self._key
 
         self.async_on_remove(
-            async_dispatcher_connect(
-                hass,
-                f"esphome_{self._entry_id}_remove_{component_key}_{key}",
+            entry_data.async_register_key_static_info_remove_callback(
+                self._static_info,
                 functools.partial(self.async_remove, force_remove=True),
             )
         )
@@ -201,10 +188,8 @@ async def async_added_to_hass(self) -> None:
             )
         )
         self.async_on_remove(
-            async_dispatcher_connect(
-                hass,
-                entry_data.signal_component_key_static_info_updated(component_key, key),
-                self._on_static_info_update,
+            entry_data.async_register_key_static_info_updated_callback(
+                self._static_info, self._on_static_info_update
             )
         )
         self._update_state_from_entry_data()