From 7655b84494f304fe48f0bb0fd93ab9cfb12c9ff8 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Fri, 1 Jul 2022 00:19:40 -0500 Subject: [PATCH] Fix key collision between platforms in esphome state updates (#74273) --- homeassistant/components/esphome/__init__.py | 19 +--- .../components/esphome/entry_data.py | 92 ++++++++++++++----- 2 files changed, 73 insertions(+), 38 deletions(-) diff --git a/homeassistant/components/esphome/__init__.py b/homeassistant/components/esphome/__init__.py index 2e88a883dc1..0c1eac3aa45 100644 --- a/homeassistant/components/esphome/__init__.py +++ b/homeassistant/components/esphome/__init__.py @@ -150,11 +150,6 @@ async def async_setup_entry( # noqa: C901 hass.bus.async_listen(EVENT_HOMEASSISTANT_STOP, on_stop) ) - @callback - def async_on_state(state: EntityState) -> None: - """Send dispatcher updates when a new state is received.""" - entry_data.async_update_state(hass, state) - @callback def async_on_service_call(service: HomeassistantServiceCall) -> None: """Call service when user automation in ESPHome config is triggered.""" @@ -288,7 +283,7 @@ async def async_setup_entry( # noqa: C901 entity_infos, services = await cli.list_entities_services() await entry_data.async_update_static_infos(hass, entry, entity_infos) await _setup_services(hass, entry_data, services) - await cli.subscribe_states(async_on_state) + await cli.subscribe_states(entry_data.async_update_state) await cli.subscribe_service_calls(async_on_service_call) await cli.subscribe_home_assistant_states(async_on_state_subscription) @@ -568,7 +563,6 @@ async def platform_async_setup_entry( @callback def async_list_entities(infos: list[EntityInfo]) -> None: """Update entities of this platform when entities are listed.""" - key_to_component = entry_data.key_to_component old_infos = entry_data.info[component_key] new_infos: dict[int, EntityInfo] = {} add_entities = [] @@ -587,12 +581,10 @@ async def platform_async_setup_entry( entity = entity_type(entry_data, component_key, info.key) add_entities.append(entity) new_infos[info.key] = info - key_to_component[info.key] = component_key # Remove old entities for info in old_infos.values(): entry_data.async_remove_entity(hass, component_key, info.key) - key_to_component.pop(info.key, None) # First copy the now-old info into the backup object entry_data.old_info[component_key] = entry_data.info[component_key] @@ -714,13 +706,8 @@ class EsphomeEntity(Entity, Generic[_InfoT, _StateT]): ) 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, + self._entry_data.async_subscribe_state_update( + self._component_key, self._key, self._on_state_update ) ) diff --git a/homeassistant/components/esphome/entry_data.py b/homeassistant/components/esphome/entry_data.py index d4bcc67db4a..8eb56e6fdb6 100644 --- a/homeassistant/components/esphome/entry_data.py +++ b/homeassistant/components/esphome/entry_data.py @@ -12,26 +12,40 @@ from aioesphomeapi import ( APIClient, APIVersion, BinarySensorInfo, + BinarySensorState, CameraInfo, + CameraState, ClimateInfo, + ClimateState, CoverInfo, + CoverState, DeviceInfo, EntityInfo, EntityState, FanInfo, + FanState, LightInfo, + LightState, LockInfo, + LockState, MediaPlayerInfo, + MediaPlayerState, NumberInfo, + NumberState, SelectInfo, + SelectState, SensorInfo, + SensorState, SwitchInfo, + SwitchState, TextSensorInfo, + TextSensorState, UserService, ) from aioesphomeapi.model import ButtonInfo from homeassistant.config_entries import ConfigEntry +from homeassistant.const import Platform from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_send from homeassistant.helpers.storage import Store @@ -41,20 +55,37 @@ _LOGGER = logging.getLogger(__name__) # Mapping from ESPHome info type to HA platform INFO_TYPE_TO_PLATFORM: dict[type[EntityInfo], str] = { - BinarySensorInfo: "binary_sensor", - ButtonInfo: "button", - CameraInfo: "camera", - ClimateInfo: "climate", - CoverInfo: "cover", - FanInfo: "fan", - LightInfo: "light", - LockInfo: "lock", - MediaPlayerInfo: "media_player", - NumberInfo: "number", - SelectInfo: "select", - SensorInfo: "sensor", - SwitchInfo: "switch", - TextSensorInfo: "sensor", + BinarySensorInfo: Platform.BINARY_SENSOR, + ButtonInfo: Platform.BINARY_SENSOR, + CameraInfo: Platform.BINARY_SENSOR, + ClimateInfo: Platform.CLIMATE, + CoverInfo: Platform.COVER, + FanInfo: Platform.FAN, + LightInfo: Platform.LIGHT, + LockInfo: Platform.LOCK, + MediaPlayerInfo: Platform.MEDIA_PLAYER, + NumberInfo: Platform.NUMBER, + SelectInfo: Platform.SELECT, + SensorInfo: Platform.SENSOR, + SwitchInfo: Platform.SWITCH, + TextSensorInfo: Platform.SENSOR, +} + +STATE_TYPE_TO_COMPONENT_KEY = { + BinarySensorState: Platform.BINARY_SENSOR, + EntityState: Platform.BINARY_SENSOR, + CameraState: Platform.BINARY_SENSOR, + ClimateState: Platform.CLIMATE, + CoverState: Platform.COVER, + FanState: Platform.FAN, + LightState: Platform.LIGHT, + LockState: Platform.LOCK, + MediaPlayerState: Platform.MEDIA_PLAYER, + NumberState: Platform.NUMBER, + SelectState: Platform.SELECT, + SensorState: Platform.SENSOR, + SwitchState: Platform.SWITCH, + TextSensorState: Platform.SENSOR, } @@ -67,7 +98,6 @@ class RuntimeEntryData: store: Store state: dict[str, dict[int, EntityState]] = field(default_factory=dict) info: dict[str, dict[int, EntityInfo]] = field(default_factory=dict) - key_to_component: dict[int, str] = field(default_factory=dict) # A second list of EntityInfo objects # This is necessary for when an entity is being removed. HA requires @@ -81,6 +111,9 @@ class RuntimeEntryData: api_version: APIVersion = field(default_factory=APIVersion) cleanup_callbacks: list[Callable[[], None]] = field(default_factory=list) disconnect_callbacks: list[Callable[[], None]] = field(default_factory=list) + state_subscriptions: dict[tuple[str, int], Callable[[], None]] = field( + default_factory=dict + ) loaded_platforms: set[str] = field(default_factory=set) platform_load_lock: asyncio.Lock = field(default_factory=asyncio.Lock) _storage_contents: dict[str, Any] | None = None @@ -125,18 +158,33 @@ class RuntimeEntryData: async_dispatcher_send(hass, signal, infos) @callback - def async_update_state(self, hass: HomeAssistant, state: EntityState) -> None: + def async_subscribe_state_update( + self, + component_key: str, + state_key: int, + entity_callback: Callable[[], None], + ) -> Callable[[], None]: + """Subscribe to state updates.""" + + def _unsubscribe() -> None: + self.state_subscriptions.pop((component_key, state_key)) + + self.state_subscriptions[(component_key, state_key)] = entity_callback + return _unsubscribe + + @callback + def async_update_state(self, state: EntityState) -> None: """Distribute an update of state information to the target.""" - component_key = self.key_to_component[state.key] + component_key = STATE_TYPE_TO_COMPONENT_KEY[type(state)] + subscription_key = (component_key, state.key) self.state[component_key][state.key] = state - signal = f"esphome_{self.entry_id}_update_{component_key}_{state.key}" _LOGGER.debug( - "Dispatching update for component %s with state key %s: %s", - component_key, - state.key, + "Dispatching update with key %s: %s", + subscription_key, state, ) - async_dispatcher_send(hass, signal) + if subscription_key in self.state_subscriptions: + self.state_subscriptions[subscription_key]() @callback def async_update_device_state(self, hass: HomeAssistant) -> None: