From 55d83140935ca9c8f26dc3635413daf8bf430a6d Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Tue, 8 Feb 2022 21:49:25 +0100
Subject: [PATCH] Fix cleanup of MQTT debug info (#66104)

---
 homeassistant/components/mqtt/__init__.py   |  2 ++
 homeassistant/components/mqtt/debug_info.py | 28 +++++++++------------
 homeassistant/components/mqtt/mixins.py     |  2 +-
 tests/components/mqtt/test_init.py          |  2 +-
 4 files changed, 16 insertions(+), 18 deletions(-)

diff --git a/homeassistant/components/mqtt/__init__.py b/homeassistant/components/mqtt/__init__.py
index 964f19e1d7c..eba2670aa95 100644
--- a/homeassistant/components/mqtt/__init__.py
+++ b/homeassistant/components/mqtt/__init__.py
@@ -596,6 +596,8 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
             )
         )
 
+    debug_info.initialize(hass)
+
     return True
 
 
diff --git a/homeassistant/components/mqtt/debug_info.py b/homeassistant/components/mqtt/debug_info.py
index e1001ab7b04..ca9e56f8efb 100644
--- a/homeassistant/components/mqtt/debug_info.py
+++ b/homeassistant/components/mqtt/debug_info.py
@@ -20,6 +20,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]:
@@ -67,9 +72,7 @@ def log_message(
     retain: bool,
 ) -> None:
     """Log an outgoing MQTT message."""
-    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": {}, "transmitted": {}}
     )
@@ -86,9 +89,7 @@ def log_message(
 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": {}, "transmitted": {}}
         )
@@ -117,9 +118,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": {}, "transmitted": {}}
     )
@@ -134,14 +133,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,
@@ -167,9 +165,7 @@ 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
diff --git a/homeassistant/components/mqtt/mixins.py b/homeassistant/components/mqtt/mixins.py
index 722bfd51c9f..177c7d9b57f 100644
--- a/homeassistant/components/mqtt/mixins.py
+++ b/homeassistant/components/mqtt/mixins.py
@@ -573,7 +573,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
 
@@ -709,6 +708,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)
 
     async def async_publish(
         self,
diff --git a/tests/components/mqtt/test_init.py b/tests/components/mqtt/test_init.py
index 3cb49598e8c..3d249fa2144 100644
--- a/tests/components/mqtt/test_init.py
+++ b/tests/components/mqtt/test_init.py
@@ -1766,7 +1766,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}")