From 86f244e57ac4dde1d9df7175f5ea70f530297670 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Wed, 30 Sep 2026 18:51:21 +0100 Subject: [PATCH] Adjust mqtt to not override Entity.add_to_platform_xxx methods (#183776) --- homeassistant/components/mqtt/entity.py | 96 ++++++++++++++-------- tests/components/mqtt/test_discovery.py | 105 +++++++++++++++++++++++- 2 files changed, 166 insertions(+), 35 deletions(-) diff --git a/homeassistant/components/mqtt/entity.py b/homeassistant/components/mqtt/entity.py index d588f421c8b1..181739d9c6b8 100644 --- a/homeassistant/components/mqtt/entity.py +++ b/homeassistant/components/mqtt/entity.py @@ -43,7 +43,10 @@ from homeassistant.helpers.dispatcher import ( async_dispatcher_send, ) from homeassistant.helpers.entity import Entity, async_generate_entity_id -from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback +from homeassistant.helpers.entity_platform import ( + AddConfigEntryEntitiesCallback, + async_get_current_platform, +) from homeassistant.helpers.event import ( async_track_device_registry_updated_event, async_track_entity_registry_updated_event, @@ -257,6 +260,22 @@ def async_setup_entity_entry_helper( # noqa: C901 ) -> None: """Set up entity creation dynamically through MQTT discovery.""" mqtt_data = hass.data[DATA_MQTT] + platform = async_get_current_platform() + + async def _async_add_discovered_entity( + entity: MqttEntity, discovery_data: DiscoveryInfoType + ) -> None: + """Add a discovered entity and acknowledge the discovery once done. + + The discovery is acknowledged only after async_add_entities returns, i.e. + after the entity wrote its initial state on success or the add was + aborted. This keeps a queued update for the same discovery hash from + draining into an entity that is not yet in the state machine. + """ + try: + await platform.async_add_entities([entity]) + finally: + send_discovery_done(hass, discovery_data) @callback def _async_migrate_subentry( @@ -335,12 +354,16 @@ def async_setup_entity_entry_helper( # noqa: C901 "and repair flow must be completed first" ) else: - async_add_entities( - [ - entity_class( - hass, config, entry, discovery_payload.discovery_data - ) - ] + entity = entity_class( + hass, config, entry, discovery_payload.discovery_data + ) + entry.async_create_task( + hass, + _async_add_discovered_entity( + entity, discovery_payload.discovery_data + ), + f"mqtt add discovered {domain} entity", + eager_start=True, ) except probatio.Invalid as err: _handle_discovery_failure(hass, discovery_payload) @@ -1010,6 +1033,7 @@ class MqttDiscoveryUpdateMixin(Entity): self._discovery_update = discovery_update self._remove_discovery_updated: Callable[[], None] | None = None self._removed_from_hass = False + self._added_to_hass = False if discovery_data is None: return mqtt_data = hass.data[DATA_MQTT] @@ -1019,6 +1043,14 @@ class MqttDiscoveryUpdateMixin(Entity): if discovery_hash in self._registry_hooks: self._registry_hooks.pop(discovery_hash)() + @override + async def async_prepare_to_add_to_hass(self) -> None: + """Register discovery cleanup that must also run if the add is aborted.""" + self._added_to_hass = False + await super().async_prepare_to_add_to_hass() + if self._discovery_data is not None: + self.async_on_remove(self._async_teardown_discovery_on_remove) + @override async def async_added_to_hass(self) -> None: """Subscribe to discovery updates.""" @@ -1219,37 +1251,32 @@ class MqttDiscoveryUpdateMixin(Entity): # rediscovered after a restart await async_remove_discovery_payload(self.hass, self._discovery_data) - @final - @override - async def add_to_platform_finish(self) -> None: - """Finish adding entity to platform.""" - await super().add_to_platform_finish() - # Only send the discovery done after the entity is fully added - # and the state is written to the state machine. - if self._discovery_data is not None: - send_discovery_done(self.hass, self._discovery_data) - @callback - @override - def add_to_platform_abort(self) -> None: - """Abort adding an entity to a platform.""" - if self._discovery_data is not None: + def _async_teardown_discovery_on_remove(self) -> None: + """Tear down discovery when the entity is removed or its add is aborted. + + Registered via async_on_remove in async_prepare_to_add_to_hass so it also + runs on the abort path, where async_will_remove_from_hass is never called. + When the add is aborted while the entity is already registered (e.g. a + disabled entity), a registry hook is installed so the retained discovery + topic is cleared if the entity is later removed from the registry. + """ + if self._discovery_data is None: + return + if not self._added_to_hass and self.registry_entry is not None: discovery_hash: tuple[str, str] = self._discovery_data[ATTR_DISCOVERY_HASH] - if self.registry_entry is not None: - self._registry_hooks[discovery_hash] = ( - async_track_entity_registry_updated_event( + self._registry_hooks[discovery_hash] = ( + async_track_entity_registry_updated_event( + self.hass, + self.entity_id, + partial( + async_clear_discovery_topic_if_entity_removed, self.hass, - self.entity_id, - partial( - async_clear_discovery_topic_if_entity_removed, - self.hass, - self._discovery_data, - ), - ) + self._discovery_data, + ), ) - stop_discovery_updates(self.hass, self._discovery_data) - send_discovery_done(self.hass, self._discovery_data) - super().add_to_platform_abort() + ) + self._cleanup_discovery_on_remove() @override async def async_will_remove_from_hass(self) -> None: @@ -1570,6 +1597,7 @@ class MqttEntity( ) await self._subscribe_topics() await self.mqtt_async_added_to_hass() + self._added_to_hass = True async def mqtt_async_added_to_hass(self) -> None: """Call before the discovery message is acknowledged. diff --git a/tests/components/mqtt/test_discovery.py b/tests/components/mqtt/test_discovery.py index b3df68d6b9a3..2258c0844600 100644 --- a/tests/components/mqtt/test_discovery.py +++ b/tests/components/mqtt/test_discovery.py @@ -26,7 +26,7 @@ from homeassistant.components.mqtt.discovery import ( MQTTDiscoveryPayload, async_start, ) -from homeassistant.components.mqtt.entity import async_removed_from_device +from homeassistant.components.mqtt.entity import MqttEntity, async_removed_from_device from homeassistant.components.mqtt.models import ReceiveMessage from homeassistant.components.mqtt.schemas import ( DEVICE_DISCOVERY_SCHEMA, @@ -1701,6 +1701,66 @@ async def test_rapid_reconfigure( assert events[2].data["new_state"].attributes["friendly_name"] == "Wine" +async def test_discovery_update_queued_until_initial_state( + hass: HomeAssistant, mqtt_mock_entry: MqttMockHAClientGenerator +) -> None: + """Test a queued discovery update is applied only after the initial state exists. + + The discovery is acknowledged caller-side, after async_add_entities returns, + so an update queued for the same discovery hash while the entity add is still + in progress must not drain into an entity that has no state yet. + """ + await mqtt_mock_entry() + events = async_capture_events(hass, EVENT_STATE_CHANGED) + + add_started = asyncio.Event() + allow_add = asyncio.Event() + original_async_added_to_hass = MqttEntity.async_added_to_hass + + async def _blocked_async_added_to_hass(self: MqttEntity) -> None: + add_started.set() + await allow_add.wait() + await original_async_added_to_hass(self) + + with patch.object(MqttEntity, "async_added_to_hass", _blocked_async_added_to_hass): + async_fire_mqtt_message( + hass, + "homeassistant/binary_sensor/bla/config", + '{ "name": "Beer", "state_topic": "test-topic" }', + ) + # Wait until the first entity add is blocked before its state is written + await add_started.wait() + + # A second payload for the same discovery hash is queued while the add + # is still in progress + async_fire_mqtt_message( + hass, + "homeassistant/binary_sensor/bla/config", + '{ "name": "Milk", "state_topic": "test-topic" }', + ) + + # The initial state does not exist yet and the queued update is not applied + assert hass.states.get("binary_sensor.beer") is None + assert not events + + allow_add.set() + await hass.async_block_till_done() + + assert len(hass.states.async_entity_ids("binary_sensor")) == 1 + state = hass.states.get("binary_sensor.beer") + assert state is not None + assert state.attributes["friendly_name"] == "Milk" + + # The initial state was written first, then the queued update was applied + assert len(events) == 2 + assert events[0].data["entity_id"] == "binary_sensor.beer" + assert events[0].data["old_state"] is None + assert events[0].data["new_state"].attributes["friendly_name"] == "Beer" + assert events[1].data["entity_id"] == "binary_sensor.beer" + assert events[1].data["old_state"].attributes["friendly_name"] == "Beer" + assert events[1].data["new_state"].attributes["friendly_name"] == "Milk" + + async def test_duplicate_removal( hass: HomeAssistant, mqtt_mock_entry: MqttMockHAClientGenerator, @@ -3042,6 +3102,49 @@ async def test_clean_up_registry_monitoring( assert len(hooks) == 0 +async def test_registry_hook_installed_when_readd_after_rename_aborts( + hass: HomeAssistant, + mqtt_mock_entry: MqttMockHAClientGenerator, + entity_registry: er.EntityRegistry, +) -> None: + """Test the registry cleanup hook is installed when an aborted re-add follows a rename. + + Renaming an entity_id makes core remove and re-add the same entity object. + _added_to_hass is set on a successful add and must be reset on every add + attempt, otherwise an aborted re-add would see the stale value and skip + installing the registry hook while the registry entry still exists, leaking + the retained discovery topic when the entity is later removed. + """ + await mqtt_mock_entry() + hooks: dict = hass.data["mqtt"].discovery_registry_hooks + config = { + "name": "milk", + "state_topic": "test-topic", + "unique_id": "very_unique", + } + async_fire_mqtt_message(hass, "homeassistant/sensor/bla/config", json.dumps(config)) + await hass.async_block_till_done() + assert hass.states.get("sensor.milk") is not None + assert len(hooks) == 0 + + async def _raise_on_readd(self: MqttEntity) -> None: + raise ValueError("Simulated re-add failure") + + # Renaming the entity_id triggers a remove and re-add of the same object; + # the patched hook aborts the re-add. + with patch.object(MqttEntity, "async_added_to_hass", _raise_on_readd): + entity_registry.async_update_entity( + "sensor.milk", new_entity_id="sensor.renamed_milk" + ) + await hass.async_block_till_done() + + # The registry entry survives the aborted re-add, so its retained discovery + # topic must be monitored for cleanup. + assert entity_registry.async_get("sensor.renamed_milk") is not None + assert len(hooks) == 1 + assert ("sensor", "bla") in hooks + + async def test_unique_id_collission_has_priority( hass: HomeAssistant, mqtt_mock_entry: MqttMockHAClientGenerator,