Clear MQTT discovery topic when a disabled entity is removed (#77757)

* Cleanup discovery on entity removal

* Add test

* Cleanup and test

* Test with clearing payload not unique id

* Address comments

* Tests cover and typing

* Just pass hass

* reuse code

* Follow up comments revert changes to cover tests

* Add test unique_id has priority over disabled

* Update homeassistant/components/mqtt/__init__.py

Co-authored-by: Erik Montnemery <erik@montnemery.com>

Co-authored-by: Erik Montnemery <erik@montnemery.com>

homeassistant/components/mqtt/__init__.py

@@ -20,7 +20,13 @@ from homeassistant.const import (
     CONF_USERNAME,
     SERVICE_RELOAD,
 )
-from homeassistant.core import HassJob, HomeAssistant, ServiceCall, callback
+from homeassistant.core import (
+    CALLBACK_TYPE,
+    HassJob,
+    HomeAssistant,
+    ServiceCall,
+    callback,
+)
 from homeassistant.exceptions import TemplateError, Unauthorized
 from homeassistant.helpers import (
     config_validation as cv,
@@ -68,6 +74,7 @@ from .const import (  # noqa: F401
     CONFIG_ENTRY_IS_SETUP,
     DATA_MQTT,
     DATA_MQTT_CONFIG,
+    DATA_MQTT_DISCOVERY_REGISTRY_HOOKS,
     DATA_MQTT_RELOAD_DISPATCHERS,
     DATA_MQTT_RELOAD_ENTRY,
     DATA_MQTT_RELOAD_NEEDED,
@@ -315,6 +322,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         # Bail out
         return False
 
+    hass.data[DATA_MQTT_DISCOVERY_REGISTRY_HOOKS] = {}
     hass.data[DATA_MQTT] = MQTT(hass, entry, conf)
     # Restore saved subscriptions
     if DATA_MQTT_SUBSCRIPTIONS_TO_RESTORE in hass.data:
@@ -638,6 +646,12 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         hass.data[DATA_MQTT_RELOAD_ENTRY] = True
     # Reload the legacy yaml platform to make entities unavailable
     await async_reload_integration_platforms(hass, DOMAIN, RELOADABLE_PLATFORMS)
+    # Cleanup entity registry hooks
+    registry_hooks: dict[tuple, CALLBACK_TYPE] = hass.data[
+        DATA_MQTT_DISCOVERY_REGISTRY_HOOKS
+    ]
+    while registry_hooks:
+        registry_hooks.popitem()[1]()
     # Wait for all ACKs and stop the loop
     await mqtt_client.async_disconnect()
     # Store remaining subscriptions to be able to restore or reload them

homeassistant/components/mqtt/const.py

@@ -33,6 +33,7 @@ CONF_TLS_VERSION = "tls_version"
 CONFIG_ENTRY_IS_SETUP = "mqtt_config_entry_is_setup"
 DATA_MQTT = "mqtt"
 DATA_MQTT_SUBSCRIPTIONS_TO_RESTORE = "mqtt_client_subscriptions"
+DATA_MQTT_DISCOVERY_REGISTRY_HOOKS = "mqtt_discovery_registry_hooks"
 DATA_MQTT_CONFIG = "mqtt_config"
 MQTT_DATA_DEVICE_TRACKER_LEGACY = "mqtt_device_tracker_legacy"
 DATA_MQTT_RELOAD_DISPATCHERS = "mqtt_reload_dispatchers"

homeassistant/components/mqtt/mixins.py

@@ -28,7 +28,13 @@ from homeassistant.const import (
     CONF_UNIQUE_ID,
     CONF_VALUE_TEMPLATE,
 )
-from homeassistant.core import Event, HomeAssistant, async_get_hass, callback
+from homeassistant.core import (
+    CALLBACK_TYPE,
+    Event,
+    HomeAssistant,
+    async_get_hass,
+    callback,
+)
 from homeassistant.helpers import (
     config_validation as cv,
     device_registry as dr,
@@ -48,6 +54,7 @@ from homeassistant.helpers.entity import (
     async_generate_entity_id,
 )
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.event import async_track_entity_registry_updated_event
 from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue
 from homeassistant.helpers.json import json_loads
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
@@ -64,6 +71,7 @@ from .const import (
     CONF_TOPIC,
     DATA_MQTT,
     DATA_MQTT_CONFIG,
+    DATA_MQTT_DISCOVERY_REGISTRY_HOOKS,
     DATA_MQTT_RELOAD_DISPATCHERS,
     DATA_MQTT_RELOAD_ENTRY,
     DATA_MQTT_UPDATED_CONFIG,
@@ -654,6 +662,17 @@ async def async_remove_discovery_payload(hass: HomeAssistant, discovery_data: di
     await async_publish(hass, discovery_topic, "", retain=True)
 
 
+async def async_clear_discovery_topic_if_entity_removed(
+    hass: HomeAssistant,
+    discovery_data: dict[str, Any],
+    event: Event,
+) -> None:
+    """Clear the discovery topic if the entity is removed."""
+    if event.data["action"] == "remove":
+        # publish empty payload to config topic to avoid re-adding
+        await async_remove_discovery_payload(hass, discovery_data)
+
+
 class MqttDiscoveryDeviceUpdate:
     """Add support for auto discovery for platforms without an entity."""
 
@@ -787,7 +806,8 @@ class MqttDiscoveryUpdate(Entity):
 
     def __init__(
         self,
-        discovery_data: dict,
+        hass: HomeAssistant,
+        discovery_data: dict | None,
         discovery_update: Callable | None = None,
     ) -> None:
         """Initialize the discovery update mixin."""
@@ -795,6 +815,14 @@ class MqttDiscoveryUpdate(Entity):
         self._discovery_update = discovery_update
         self._remove_discovery_updated: Callable | None = None
         self._removed_from_hass = False
+        if discovery_data is None:
+            return
+        self._registry_hooks: dict[tuple, CALLBACK_TYPE] = hass.data[
+            DATA_MQTT_DISCOVERY_REGISTRY_HOOKS
+        ]
+        discovery_hash: tuple[str, str] = discovery_data[ATTR_DISCOVERY_HASH]
+        if discovery_hash in self._registry_hooks:
+            self._registry_hooks.pop(discovery_hash)()
 
     async def async_added_to_hass(self) -> None:
         """Subscribe to discovery updates."""
@@ -857,7 +885,7 @@ class MqttDiscoveryUpdate(Entity):
 
     async def async_removed_from_registry(self) -> None:
         """Clear retained discovery topic in broker."""
-        if not self._removed_from_hass:
+        if not self._removed_from_hass and self._discovery_data is not None:
             # Stop subscribing to discovery updates to not trigger when we clear the
             # discovery topic
             self._cleanup_discovery_on_remove()
@@ -868,7 +896,20 @@ class MqttDiscoveryUpdate(Entity):
     @callback
     def add_to_platform_abort(self) -> None:
         """Abort adding an entity to a platform."""
-        if self._discovery_data:
+        if self._discovery_data is not None:
+            discovery_hash: tuple = 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.hass,
+                    self.entity_id,
+                    partial(
+                        async_clear_discovery_topic_if_entity_removed,
+                        self.hass,
+                        self._discovery_data,
+                    ),
+                )
             stop_discovery_updates(self.hass, self._discovery_data)
             send_discovery_done(self.hass, self._discovery_data)
         super().add_to_platform_abort()
@@ -976,7 +1017,7 @@ class MqttEntity(
         # Initialize mixin classes
         MqttAttributes.__init__(self, config)
         MqttAvailability.__init__(self, config)
-        MqttDiscoveryUpdate.__init__(self, discovery_data, self.discovery_update)
+        MqttDiscoveryUpdate.__init__(self, hass, discovery_data, self.discovery_update)
         MqttEntityDeviceInfo.__init__(self, config.get(CONF_DEVICE), config_entry)
 
     def _init_entity_id(self):

tests/components/mqtt/test_discovery.py

@@ -1,4 +1,5 @@
 """The tests for the MQTT discovery."""
+import copy
 import json
 from pathlib import Path
 import re
@@ -23,6 +24,8 @@ from homeassistant.const import (
 import homeassistant.core as ha
 from homeassistant.setup import async_setup_component
 
+from .test_common import help_test_unload_config_entry
+
 from tests.common import (
     MockConfigEntry,
     async_capture_events,
@@ -1356,3 +1359,170 @@ async def test_mqtt_discovery_unsubscribe_once(
         await hass.async_block_till_done()
         await hass.async_block_till_done()
         mqtt_client_mock.unsubscribe.assert_called_once_with("comp/discovery/#")
+
+
+@patch("homeassistant.components.mqtt.PLATFORMS", [Platform.SENSOR])
+async def test_clear_config_topic_disabled_entity(
+    hass, mqtt_mock_entry_no_yaml_config, device_reg, caplog
+):
+    """Test the discovery topic is removed when a disabled entity is removed."""
+    mqtt_mock = await mqtt_mock_entry_no_yaml_config()
+    # discover an entity that is not enabled by default
+    config = {
+        "name": "sbfspot_12345",
+        "state_topic": "homeassistant_test/sensor/sbfspot_0/sbfspot_12345/",
+        "unique_id": "sbfspot_12345",
+        "enabled_by_default": False,
+        "device": {
+            "identifiers": ["sbfspot_12345"],
+            "name": "sbfspot_12345",
+            "sw_version": "1.0",
+            "connections": [["mac", "12:34:56:AB:CD:EF"]],
+        },
+    }
+    async_fire_mqtt_message(
+        hass,
+        "homeassistant/sensor/sbfspot_0/sbfspot_12345/config",
+        json.dumps(config),
+    )
+    await hass.async_block_till_done()
+    # discover an entity that is not unique (part 1), will be added
+    config_not_unique1 = copy.deepcopy(config)
+    config_not_unique1["name"] = "sbfspot_12345_1"
+    config_not_unique1["unique_id"] = "not_unique"
+    config_not_unique1.pop("enabled_by_default")
+    async_fire_mqtt_message(
+        hass,
+        "homeassistant/sensor/sbfspot_0/sbfspot_12345_1/config",
+        json.dumps(config_not_unique1),
+    )
+    # discover an entity that is not unique (part 2), will not be added
+    config_not_unique2 = copy.deepcopy(config_not_unique1)
+    config_not_unique2["name"] = "sbfspot_12345_2"
+    async_fire_mqtt_message(
+        hass,
+        "homeassistant/sensor/sbfspot_0/sbfspot_12345_2/config",
+        json.dumps(config_not_unique2),
+    )
+    await hass.async_block_till_done()
+    assert "Platform mqtt does not generate unique IDs" in caplog.text
+
+    assert hass.states.get("sensor.sbfspot_12345") is None  # disabled
+    assert hass.states.get("sensor.sbfspot_12345_1") is not None  # enabled
+    assert hass.states.get("sensor.sbfspot_12345_2") is None  # not unique
+
+    # Verify device is created
+    device_entry = device_reg.async_get_device(set(), {("mac", "12:34:56:AB:CD:EF")})
+    assert device_entry is not None
+
+    # Remove the device from the registry
+    device_reg.async_remove_device(device_entry.id)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    # Assert all valid discovery topics are cleared
+    assert mqtt_mock.async_publish.call_count == 2
+    assert (
+        call("homeassistant/sensor/sbfspot_0/sbfspot_12345/config", "", 0, True)
+        in mqtt_mock.async_publish.mock_calls
+    )
+    assert (
+        call("homeassistant/sensor/sbfspot_0/sbfspot_12345_1/config", "", 0, True)
+        in mqtt_mock.async_publish.mock_calls
+    )
+
+
+@patch("homeassistant.components.mqtt.PLATFORMS", [Platform.SENSOR])
+async def test_clean_up_registry_monitoring(
+    hass, mqtt_mock_entry_no_yaml_config, device_reg, tmp_path
+):
+    """Test registry monitoring hook is removed after a reload."""
+    await mqtt_mock_entry_no_yaml_config()
+    hooks: dict = hass.data[mqtt.const.DATA_MQTT_DISCOVERY_REGISTRY_HOOKS]
+    # discover an entity that is not enabled by default
+    config1 = {
+        "name": "sbfspot_12345",
+        "state_topic": "homeassistant_test/sensor/sbfspot_0/sbfspot_12345/",
+        "unique_id": "sbfspot_12345",
+        "enabled_by_default": False,
+        "device": {
+            "identifiers": ["sbfspot_12345"],
+            "name": "sbfspot_12345",
+            "sw_version": "1.0",
+            "connections": [["mac", "12:34:56:AB:CD:EF"]],
+        },
+    }
+    # Publish it config
+    # Since it is not enabled_by_default the sensor will not be loaded
+    # it should register a hook for monitoring the entiry registry
+    async_fire_mqtt_message(
+        hass,
+        "homeassistant/sensor/sbfspot_0/sbfspot_12345/config",
+        json.dumps(config1),
+    )
+    await hass.async_block_till_done()
+    assert len(hooks) == 1
+
+    # Publish it again no new monitor should be started
+    async_fire_mqtt_message(
+        hass,
+        "homeassistant/sensor/sbfspot_0/sbfspot_12345/config",
+        json.dumps(config1),
+    )
+    await hass.async_block_till_done()
+    assert len(hooks) == 1
+
+    # Verify device is created
+    device_entry = device_reg.async_get_device(set(), {("mac", "12:34:56:AB:CD:EF")})
+    assert device_entry is not None
+
+    # Enload the entry
+    # The monitoring should be cleared
+    await help_test_unload_config_entry(hass, tmp_path, {})
+    assert len(hooks) == 0
+
+
+@patch("homeassistant.components.mqtt.PLATFORMS", [Platform.SENSOR])
+async def test_unique_id_collission_has_priority(
+    hass, mqtt_mock_entry_no_yaml_config, entity_reg
+):
+    """Test tehe unique_id collision detection has priority over registry disabled items."""
+    await mqtt_mock_entry_no_yaml_config()
+    config = {
+        "name": "sbfspot_12345",
+        "state_topic": "homeassistant_test/sensor/sbfspot_0/sbfspot_12345/",
+        "unique_id": "sbfspot_12345",
+        "enabled_by_default": False,
+        "device": {
+            "identifiers": ["sbfspot_12345"],
+            "name": "sbfspot_12345",
+            "sw_version": "1.0",
+            "connections": [["mac", "12:34:56:AB:CD:EF"]],
+        },
+    }
+    # discover an entity that is not unique and disabled by default (part 1), will be added
+    config_not_unique1 = copy.deepcopy(config)
+    config_not_unique1["name"] = "sbfspot_12345_1"
+    config_not_unique1["unique_id"] = "not_unique"
+    async_fire_mqtt_message(
+        hass,
+        "homeassistant/sensor/sbfspot_0/sbfspot_12345_1/config",
+        json.dumps(config_not_unique1),
+    )
+    # discover an entity that is not unique (part 2), will not be added, and the registry entry is cleared
+    config_not_unique2 = copy.deepcopy(config_not_unique1)
+    config_not_unique2["name"] = "sbfspot_12345_2"
+    async_fire_mqtt_message(
+        hass,
+        "homeassistant/sensor/sbfspot_0/sbfspot_12345_2/config",
+        json.dumps(config_not_unique2),
+    )
+    await hass.async_block_till_done()
+
+    assert hass.states.get("sensor.sbfspot_12345_1") is None  # not enabled
+    assert hass.states.get("sensor.sbfspot_12345_2") is None  # not unique
+
+    # Verify the first entity is created
+    assert entity_reg.async_get("sensor.sbfspot_12345_1") is not None
+    # Verify the second entity is not created because it is not unique
+    assert entity_reg.async_get("sensor.sbfspot_12345_2") is None