diff --git a/homeassistant/components/esphome/__init__.py b/homeassistant/components/esphome/__init__.py index c9d07a22ec6..a12754a87f4 100644 --- a/homeassistant/components/esphome/__init__.py +++ b/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, + ) + ) diff --git a/homeassistant/components/esphome/camera.py b/homeassistant/components/esphome/camera.py index 1678281c3de..5b8f4f0d7e6 100644 --- a/homeassistant/components/esphome/camera.py +++ b/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()