Convert most esphome state updates to callbacks (#43246)

homeassistant/components/esphome/__init__.py

@@ -489,58 +489,38 @@ def esphome_map_enum(func: Callable[[], Dict[int, str]]):
     return EsphomeEnumMapper(func)
 
 
-class EsphomeEntity(Entity):
-    """Define a generic esphome entity."""
+class EsphomeBaseEntity(Entity):
+    """Define a base esphome entity."""
 
     def __init__(self, entry_id: str, component_key: str, key: int):
         """Initialize."""
         self._entry_id = entry_id
         self._component_key = component_key
         self._key = key
-        self._remove_callbacks: List[Callable[[], None]] = []
 
     async def async_added_to_hass(self) -> None:
         """Register callbacks."""
-        kwargs = {
-            "entry_id": self._entry_id,
-            "component_key": self._component_key,
-            "key": self._key,
-        }
-        self._remove_callbacks.append(
+        self.async_on_remove(
             async_dispatcher_connect(
                 self.hass,
                 (
-                    f"esphome_{kwargs.get('entry_id')}"
-                    f"_update_{kwargs.get('component_key')}_{kwargs.get('key')}"
-                ),
-                self._on_state_update,
-            )
-        )
-
-        self._remove_callbacks.append(
-            async_dispatcher_connect(
-                self.hass,
-                (
-                    f"esphome_{kwargs.get('entry_id')}_remove_"
-                    f"{kwargs.get('component_key')}_{kwargs.get('key')}"
+                    f"esphome_{self._entry_id}_remove_"
+                    f"{self._component_key}_{self._key}"
                 ),
                 self.async_remove,
             )
         )
 
-        self._remove_callbacks.append(
+        self.async_on_remove(
             async_dispatcher_connect(
                 self.hass,
-                f"esphome_{kwargs.get('entry_id')}_on_device_update",
+                f"esphome_{self._entry_id}_on_device_update",
                 self._on_device_update,
             )
         )
 
-    async def _on_state_update(self) -> None:
-        """Update the entity state when state or static info changed."""
-        self.async_write_ha_state()
-
-    async def _on_device_update(self) -> None:
+    @callback
+    def _on_device_update(self) -> None:
         """Update the entity state when device info has changed."""
         if self._entry_data.available:
             # Don't update the HA state yet when the device comes online.
@@ -549,12 +529,6 @@ class EsphomeEntity(Entity):
             return
         self.async_write_ha_state()
 
-    async def async_will_remove_from_hass(self) -> None:
-        """Unregister callbacks."""
-        for remove_callback in self._remove_callbacks:
-            remove_callback()
-        self._remove_callbacks = []
-
     @property
     def _entry_data(self) -> RuntimeEntryData:
         return self.hass.data[DATA_KEY][self._entry_id]
@@ -619,3 +593,23 @@ class EsphomeEntity(Entity):
     def should_poll(self) -> bool:
         """Disable polling."""
         return False
+
+
+class EsphomeEntity(EsphomeBaseEntity):
+    """Define a generic esphome entity."""
+
+    async def async_added_to_hass(self) -> None:
+        """Register callbacks."""
+
+        await super().async_added_to_hass()
+
+        self.async_on_remove(
+            async_dispatcher_connect(
+                self.hass,
+                (
+                    f"esphome_{self._entry_id}"
+                    f"_update_{self._component_key}_{self._key}"
+                ),
+                self.async_write_ha_state,
+            )
+        )

homeassistant/components/esphome/camera.py

@@ -7,9 +7,10 @@ from aioesphomeapi import CameraInfo, CameraState
 from homeassistant.components import camera
 from homeassistant.components.camera import Camera
 from homeassistant.config_entries import ConfigEntry
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.typing import HomeAssistantType
 
-from . import EsphomeEntity, platform_async_setup_entry
+from . import EsphomeBaseEntity, platform_async_setup_entry
 
 
 async def async_setup_entry(
@@ -27,13 +28,13 @@ async def async_setup_entry(
     )
 
 
-class EsphomeCamera(Camera, EsphomeEntity):
+class EsphomeCamera(Camera, EsphomeBaseEntity):
     """A camera implementation for ESPHome."""
 
     def __init__(self, entry_id: str, component_key: str, key: int):
         """Initialize."""
         Camera.__init__(self)
-        EsphomeEntity.__init__(self, entry_id, component_key, key)
+        EsphomeBaseEntity.__init__(self, entry_id, component_key, key)
         self._image_cond = asyncio.Condition()
 
     @property
@@ -44,9 +45,25 @@ class EsphomeCamera(Camera, EsphomeEntity):
     def _state(self) -> Optional[CameraState]:
         return super()._state
 
+    async def async_added_to_hass(self) -> None:
+        """Register callbacks."""
+
+        await super().async_added_to_hass()
+
+        self.async_on_remove(
+            async_dispatcher_connect(
+                self.hass,
+                (
+                    f"esphome_{self._entry_id}"
+                    f"_update_{self._component_key}_{self._key}"
+                ),
+                self._on_state_update,
+            )
+        )
+
     async def _on_state_update(self) -> None:
         """Notify listeners of new image when update arrives."""
-        await super()._on_state_update()
+        self.async_write_ha_state()
         async with self._image_cond:
             self._image_cond.notify_all()