Adjust mqtt to not override Entity.add_to_platform_xxx methods (#183776)

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.

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,