Fix cleanup of MQTT debug info (#66104)

homeassistant/components/mqtt/__init__.py

@@ -595,6 +595,8 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
             )
         )
 
+    debug_info.initialize(hass)
+
     return True
 
 

homeassistant/components/mqtt/debug_info.py

@@ -15,6 +15,11 @@ DATA_MQTT_DEBUG_INFO = "mqtt_debug_info"
 STORED_MESSAGES = 10
 
 
+def initialize(hass: HomeAssistant):
+    """Initialize MQTT debug info."""
+    hass.data[DATA_MQTT_DEBUG_INFO] = {"entities": {}, "triggers": {}}
+
+
 def log_messages(
     hass: HomeAssistant, entity_id: str
 ) -> Callable[[MessageCallbackType], MessageCallbackType]:
@@ -45,9 +50,7 @@ def log_messages(
 def add_subscription(hass, message_callback, subscription):
     """Prepare debug data for subscription."""
     if entity_id := getattr(message_callback, "__entity_id", None):
-        debug_info = hass.data.setdefault(
-            DATA_MQTT_DEBUG_INFO, {"entities": {}, "triggers": {}}
-        )
+        debug_info = hass.data[DATA_MQTT_DEBUG_INFO]
         entity_info = debug_info["entities"].setdefault(
             entity_id, {"subscriptions": {}, "discovery_data": {}}
         )
@@ -76,9 +79,7 @@ def remove_subscription(hass, message_callback, subscription):
 
 def add_entity_discovery_data(hass, discovery_data, entity_id):
     """Add discovery data."""
-    debug_info = hass.data.setdefault(
-        DATA_MQTT_DEBUG_INFO, {"entities": {}, "triggers": {}}
-    )
+    debug_info = hass.data[DATA_MQTT_DEBUG_INFO]
     entity_info = debug_info["entities"].setdefault(
         entity_id, {"subscriptions": {}, "discovery_data": {}}
     )
@@ -93,14 +94,13 @@ def update_entity_discovery_data(hass, discovery_payload, entity_id):
 
 def remove_entity_data(hass, entity_id):
     """Remove discovery data."""
-    hass.data[DATA_MQTT_DEBUG_INFO]["entities"].pop(entity_id)
+    if entity_id in hass.data[DATA_MQTT_DEBUG_INFO]["entities"]:
+        hass.data[DATA_MQTT_DEBUG_INFO]["entities"].pop(entity_id)
 
 
 def add_trigger_discovery_data(hass, discovery_hash, discovery_data, device_id):
     """Add discovery data."""
-    debug_info = hass.data.setdefault(
-        DATA_MQTT_DEBUG_INFO, {"entities": {}, "triggers": {}}
-    )
+    debug_info = hass.data[DATA_MQTT_DEBUG_INFO]
     debug_info["triggers"][discovery_hash] = {
         "device_id": device_id,
         "discovery_data": discovery_data,
@@ -126,9 +126,7 @@ async def info_for_device(hass, device_id):
     entries = hass.helpers.entity_registry.async_entries_for_device(
         entity_registry, device_id, include_disabled_entities=True
     )
-    mqtt_debug_info = hass.data.setdefault(
-        DATA_MQTT_DEBUG_INFO, {"entities": {}, "triggers": {}}
-    )
+    mqtt_debug_info = hass.data[DATA_MQTT_DEBUG_INFO]
     for entry in entries:
         if entry.entity_id not in mqtt_debug_info["entities"]:
             continue

homeassistant/components/mqtt/mixins.py

@@ -549,7 +549,6 @@ class MqttDiscoveryUpdate(Entity):
     def _cleanup_discovery_on_remove(self) -> None:
         """Stop listening to signal and cleanup discovery data."""
         if self._discovery_data and not self._removed_from_hass:
-            debug_info.remove_entity_data(self.hass, self.entity_id)
             clear_discovery_hash(self.hass, self._discovery_data[ATTR_DISCOVERY_HASH])
             self._removed_from_hass = True
 
@@ -677,6 +676,7 @@ class MqttEntity(
         await MqttAttributes.async_will_remove_from_hass(self)
         await MqttAvailability.async_will_remove_from_hass(self)
         await MqttDiscoveryUpdate.async_will_remove_from_hass(self)
+        debug_info.remove_entity_data(self.hass, self.entity_id)
 
     @staticmethod
     @abstractmethod

tests/components/mqtt/test_init.py

@@ -1764,7 +1764,7 @@ async def test_debug_info_multiple_entities_triggers(hass, mqtt_mock):
         } in discovery_data
 
 
-async def test_debug_info_non_mqtt(hass, device_reg, entity_reg):
+async def test_debug_info_non_mqtt(hass, device_reg, entity_reg, mqtt_mock):
     """Test we get empty debug_info for a device with non MQTT entities."""
     DOMAIN = "sensor"
     platform = getattr(hass.components, f"test.{DOMAIN}")