Fix MQTT debug info for same topic (#34952)

Author: emontnemery

homeassistant/components/mqtt/debug_info.py

@@ -23,7 +23,7 @@ def _log_message(msg):
         debug_info = hass.data[DATA_MQTT_DEBUG_INFO]
         messages = debug_info["entities"][entity_id]["subscriptions"][
             msg.subscribed_topic
-        ]
+        ]["messages"]
         if msg not in messages:
             messages.append(msg)
 
@@ -50,16 +50,27 @@ def add_subscription(hass, message_callback, subscription):
         entity_info = debug_info["entities"].setdefault(
             entity_id, {"subscriptions": {}, "discovery_data": {}}
         )
-        entity_info["subscriptions"][subscription] = deque([], STORED_MESSAGES)
+        if subscription not in entity_info["subscriptions"]:
+            entity_info["subscriptions"][subscription] = {
+                "count": 0,
+                "messages": deque([], STORED_MESSAGES),
+            }
+        entity_info["subscriptions"][subscription]["count"] += 1
 
 
 def remove_subscription(hass, message_callback, subscription):
-    """Remove debug data for subscription."""
+    """Remove debug data for subscription if it exists."""
     entity_id = getattr(message_callback, "__entity_id", None)
     if entity_id and entity_id in hass.data[DATA_MQTT_DEBUG_INFO]["entities"]:
-        hass.data[DATA_MQTT_DEBUG_INFO]["entities"][entity_id]["subscriptions"].pop(
+        hass.data[DATA_MQTT_DEBUG_INFO]["entities"][entity_id]["subscriptions"][
             subscription
-        )
+        ]["count"] -= 1
+        if not hass.data[DATA_MQTT_DEBUG_INFO]["entities"][entity_id]["subscriptions"][
+            subscription
+        ]["count"]:
+            hass.data[DATA_MQTT_DEBUG_INFO]["entities"][entity_id]["subscriptions"].pop(
+                subscription
+            )
 
 
 def add_entity_discovery_data(hass, discovery_data, entity_id):
@@ -127,10 +138,10 @@ async def info_for_device(hass, device_id):
                 "topic": topic,
                 "messages": [
                     {"payload": msg.payload, "time": msg.timestamp, "topic": msg.topic}
-                    for msg in list(messages)
+                    for msg in list(subscription["messages"])
                 ],
             }
-            for topic, messages in entity_info["subscriptions"].items()
+            for topic, subscription in entity_info["subscriptions"].items()
         ]
         discovery_data = {
             "topic": entity_info["discovery_data"].get(ATTR_DISCOVERY_TOPIC, ""),

tests/components/mqtt/test_init.py

@@ -956,6 +956,42 @@ async def test_mqtt_ws_remove_discovered_device_twice(
     assert response["error"]["code"] == websocket_api.const.ERR_NOT_FOUND
 
 
+async def test_mqtt_ws_remove_discovered_device_same_topic(
+    hass, device_reg, hass_ws_client, mqtt_mock
+):
+    """Test MQTT websocket device removal."""
+    config_entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    config_entry.add_to_hass(hass)
+    await async_start(hass, "homeassistant", {}, config_entry)
+
+    data = (
+        '{ "device":{"identifiers":["0AFFD2"]},'
+        '  "state_topic": "foobar/sensor",'
+        '  "availability_topic": "foobar/sensor",'
+        '  "unique_id": "unique" }'
+    )
+
+    async_fire_mqtt_message(hass, "homeassistant/sensor/bla/config", data)
+    await hass.async_block_till_done()
+
+    device_entry = device_reg.async_get_device({("mqtt", "0AFFD2")}, set())
+    assert device_entry is not None
+
+    client = await hass_ws_client(hass)
+    await client.send_json(
+        {"id": 5, "type": "mqtt/device/remove", "device_id": device_entry.id}
+    )
+    response = await client.receive_json()
+    assert response["success"]
+
+    await client.send_json(
+        {"id": 6, "type": "mqtt/device/remove", "device_id": device_entry.id}
+    )
+    response = await client.receive_json()
+    assert not response["success"]
+    assert response["error"]["code"] == websocket_api.const.ERR_NOT_FOUND
+
+
 async def test_mqtt_ws_remove_non_mqtt_device(
     hass, device_reg, hass_ws_client, mqtt_mock
 ):
@@ -1302,7 +1338,60 @@ async def test_debug_info_filter_same(hass, mqtt_mock):
     assert {
         "topic": "sensor/#",
         "messages": [
-            {"topic": "sensor/abc", "payload": "123", "time": dt1},
-            {"topic": "sensor/abc", "payload": "123", "time": dt2},
+            {"payload": "123", "time": dt1, "topic": "sensor/abc"},
+            {"payload": "123", "time": dt2, "topic": "sensor/abc"},
         ],
     } == debug_info_data["entities"][0]["subscriptions"][0]
+
+
+async def test_debug_info_same_topic(hass, mqtt_mock):
+    """Test debug info."""
+    config = {
+        "device": {"identifiers": ["helloworld"]},
+        "platform": "mqtt",
+        "name": "test",
+        "state_topic": "sensor/status",
+        "availability_topic": "sensor/status",
+        "unique_id": "veryunique",
+    }
+
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    entry.add_to_hass(hass)
+    await async_start(hass, "homeassistant", {}, entry)
+    registry = await hass.helpers.device_registry.async_get_registry()
+
+    data = json.dumps(config)
+    async_fire_mqtt_message(hass, "homeassistant/sensor/bla/config", data)
+    await hass.async_block_till_done()
+
+    device = registry.async_get_device({("mqtt", "helloworld")}, set())
+    assert device is not None
+
+    debug_info_data = await debug_info.info_for_device(hass, device.id)
+    assert len(debug_info_data["entities"][0]["subscriptions"]) >= 1
+    assert {"topic": "sensor/status", "messages": []} in debug_info_data["entities"][0][
+        "subscriptions"
+    ]
+
+    start_dt = datetime(2019, 1, 1, 0, 0, 0)
+    with patch("homeassistant.util.dt.utcnow") as dt_utcnow:
+        dt_utcnow.return_value = start_dt
+        async_fire_mqtt_message(hass, "sensor/status", "123", qos=0, retain=False)
+
+    debug_info_data = await debug_info.info_for_device(hass, device.id)
+    assert len(debug_info_data["entities"][0]["subscriptions"]) == 1
+    assert {
+        "payload": "123",
+        "time": start_dt,
+        "topic": "sensor/status",
+    } in debug_info_data["entities"][0]["subscriptions"][0]["messages"]
+
+    config["availability_topic"] = "sensor/availability"
+    data = json.dumps(config)
+    async_fire_mqtt_message(hass, "homeassistant/sensor/bla/config", data)
+    await hass.async_block_till_done()
+
+    start_dt = datetime(2019, 1, 1, 0, 0, 0)
+    with patch("homeassistant.util.dt.utcnow") as dt_utcnow:
+        dt_utcnow.return_value = start_dt
+        async_fire_mqtt_message(hass, "sensor/status", "123", qos=0, retain=False)