From 324dfe07b464afd4e1c66622e917cd44bea56a35 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Tue, 10 Mar 2020 05:59:38 +0100
Subject: [PATCH] Clear discovery topic for MQTT device triggers (#32617)

---
 .../components/mqtt/device_trigger.py         | 16 +++++---
 tests/components/mqtt/test_device_trigger.py  | 39 +++++++++++++++++++
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/homeassistant/components/mqtt/device_trigger.py b/homeassistant/components/mqtt/device_trigger.py
index 92bef0578c9..5bb5ccbd9d4 100644
--- a/homeassistant/components/mqtt/device_trigger.py
+++ b/homeassistant/components/mqtt/device_trigger.py
@@ -18,6 +18,7 @@ from homeassistant.helpers.typing import ConfigType, HomeAssistantType
 
 from . import (
     ATTR_DISCOVERY_HASH,
+    ATTR_DISCOVERY_TOPIC,
     CONF_CONNECTIONS,
     CONF_DEVICE,
     CONF_IDENTIFIERS,
@@ -99,7 +100,7 @@ class Trigger:
     """Device trigger settings."""
 
     device_id = attr.ib(type=str)
-    discovery_hash = attr.ib(type=tuple)
+    discovery_data = attr.ib(type=dict)
     hass = attr.ib(type=HomeAssistantType)
     payload = attr.ib(type=str)
     qos = attr.ib(type=int)
@@ -132,7 +133,6 @@ class Trigger:
 
     async def update_trigger(self, config, discovery_hash, remove_signal):
         """Update MQTT device trigger."""
-        self.discovery_hash = discovery_hash
         self.remove_signal = remove_signal
         self.type = config[CONF_TYPE]
         self.subtype = config[CONF_SUBTYPE]
@@ -216,7 +216,7 @@ async def async_setup_trigger(hass, config, config_entry, discovery_data):
         hass.data[DEVICE_TRIGGERS][discovery_id] = Trigger(
             hass=hass,
             device_id=device.id,
-            discovery_hash=discovery_hash,
+            discovery_data=discovery_data,
             type=config[CONF_TYPE],
             subtype=config[CONF_SUBTYPE],
             topic=config[CONF_TOPIC],
@@ -236,9 +236,15 @@ async def async_device_removed(hass: HomeAssistant, device_id: str):
     for trig in triggers:
         device_trigger = hass.data[DEVICE_TRIGGERS].pop(trig[CONF_DISCOVERY_ID])
         if device_trigger:
+            discovery_hash = device_trigger.discovery_data[ATTR_DISCOVERY_HASH]
+            discovery_topic = device_trigger.discovery_data[ATTR_DISCOVERY_TOPIC]
+
             device_trigger.detach_trigger()
-            clear_discovery_hash(hass, device_trigger.discovery_hash)
+            clear_discovery_hash(hass, discovery_hash)
             device_trigger.remove_signal()
+            mqtt.publish(
+                hass, discovery_topic, "", retain=True,
+            )
 
 
 async def async_get_triggers(hass: HomeAssistant, device_id: str) -> List[dict]:
@@ -281,7 +287,7 @@ async def async_attach_trigger(
         hass.data[DEVICE_TRIGGERS][discovery_id] = Trigger(
             hass=hass,
             device_id=device_id,
-            discovery_hash=None,
+            discovery_data=None,
             remove_signal=None,
             type=config[CONF_TYPE],
             subtype=config[CONF_SUBTYPE],
diff --git a/tests/components/mqtt/test_device_trigger.py b/tests/components/mqtt/test_device_trigger.py
index c9d9ec4ad08..c7d1f636c02 100644
--- a/tests/components/mqtt/test_device_trigger.py
+++ b/tests/components/mqtt/test_device_trigger.py
@@ -831,3 +831,42 @@ async def test_entity_device_info_update(hass, mqtt_mock):
     device = registry.async_get_device({("mqtt", "helloworld")}, set())
     assert device is not None
     assert device.name == "Milk"
+
+
+async def test_cleanup_device(hass, device_reg, entity_reg, mqtt_mock):
+    """Test discovered device is cleaned up when removed from registry."""
+    config_entry = MockConfigEntry(domain=DOMAIN)
+    config_entry.add_to_hass(hass)
+    await async_start(hass, "homeassistant", {}, config_entry)
+
+    config = {
+        "automation_type": "trigger",
+        "topic": "test-topic",
+        "type": "foo",
+        "subtype": "bar",
+        "device": {"identifiers": ["helloworld"]},
+    }
+
+    data = json.dumps(config)
+    async_fire_mqtt_message(hass, "homeassistant/device_automation/bla/config", data)
+    await hass.async_block_till_done()
+
+    # Verify device registry entry is created
+    device_entry = device_reg.async_get_device({("mqtt", "helloworld")}, set())
+    assert device_entry is not None
+
+    triggers = await async_get_device_automations(hass, "trigger", device_entry.id)
+    assert triggers[0]["type"] == "foo"
+
+    device_reg.async_remove_device(device_entry.id)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    # Verify device registry entry is cleared
+    device_entry = device_reg.async_get_device({("mqtt", "0AFFD2")}, set())
+    assert device_entry is None
+
+    # Verify retained discovery topic has been cleared
+    mqtt_mock.async_publish.assert_called_once_with(
+        "homeassistant/device_automation/bla/config", "", 0, True
+    )