From 42790d3e970ebe82c4a5a4358ed0ae7136f88a48 Mon Sep 17 00:00:00 2001
From: emontnemery <erik@montnemery.com>
Date: Tue, 25 Sep 2018 08:44:14 +0200
Subject: [PATCH] Remove discovered MQTT alarm_control_panel device when
 discovery topic is cleared (#16825)

---
 .../components/alarm_control_panel/mqtt.py    | 28 +++++++++++++------
 tests/components/mqtt/test_discovery.py       | 25 +++++++++++++++++
 2 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/homeassistant/components/alarm_control_panel/mqtt.py b/homeassistant/components/alarm_control_panel/mqtt.py
index e5ad54c4147..e36765b2460 100644
--- a/homeassistant/components/alarm_control_panel/mqtt.py
+++ b/homeassistant/components/alarm_control_panel/mqtt.py
@@ -18,9 +18,9 @@ from homeassistant.const import (
     STATE_ALARM_PENDING, STATE_ALARM_TRIGGERED, STATE_UNKNOWN,
     CONF_NAME, CONF_CODE)
 from homeassistant.components.mqtt import (
-    CONF_AVAILABILITY_TOPIC, CONF_STATE_TOPIC, CONF_COMMAND_TOPIC,
-    CONF_PAYLOAD_AVAILABLE, CONF_PAYLOAD_NOT_AVAILABLE, CONF_QOS,
-    CONF_RETAIN, MqttAvailability)
+    ATTR_DISCOVERY_HASH, CONF_AVAILABILITY_TOPIC, CONF_STATE_TOPIC,
+    CONF_COMMAND_TOPIC, CONF_PAYLOAD_AVAILABLE, CONF_PAYLOAD_NOT_AVAILABLE,
+    CONF_QOS, CONF_RETAIN, MqttAvailability, MqttDiscoveryUpdate)
 import homeassistant.helpers.config_validation as cv
 
 _LOGGER = logging.getLogger(__name__)
@@ -53,6 +53,10 @@ def async_setup_platform(hass, config, async_add_entities,
     if discovery_info is not None:
         config = PLATFORM_SCHEMA(discovery_info)
 
+    discovery_hash = None
+    if discovery_info is not None and ATTR_DISCOVERY_HASH in discovery_info:
+        discovery_hash = discovery_info[ATTR_DISCOVERY_HASH]
+
     async_add_entities([MqttAlarm(
         config.get(CONF_NAME),
         config.get(CONF_STATE_TOPIC),
@@ -65,18 +69,22 @@ def async_setup_platform(hass, config, async_add_entities,
         config.get(CONF_CODE),
         config.get(CONF_AVAILABILITY_TOPIC),
         config.get(CONF_PAYLOAD_AVAILABLE),
-        config.get(CONF_PAYLOAD_NOT_AVAILABLE))])
+        config.get(CONF_PAYLOAD_NOT_AVAILABLE),
+        discovery_hash,)])
 
 
-class MqttAlarm(MqttAvailability, alarm.AlarmControlPanel):
+class MqttAlarm(MqttAvailability, MqttDiscoveryUpdate,
+                alarm.AlarmControlPanel):
     """Representation of a MQTT alarm status."""
 
     def __init__(self, name, state_topic, command_topic, qos, retain,
                  payload_disarm, payload_arm_home, payload_arm_away, code,
-                 availability_topic, payload_available, payload_not_available):
+                 availability_topic, payload_available, payload_not_available,
+                 discovery_hash):
         """Init the MQTT Alarm Control Panel."""
-        super().__init__(availability_topic, qos, payload_available,
-                         payload_not_available)
+        MqttAvailability.__init__(self, availability_topic, qos,
+                                  payload_available, payload_not_available)
+        MqttDiscoveryUpdate.__init__(self, discovery_hash)
         self._state = STATE_UNKNOWN
         self._name = name
         self._state_topic = state_topic
@@ -87,11 +95,13 @@ class MqttAlarm(MqttAvailability, alarm.AlarmControlPanel):
         self._payload_arm_home = payload_arm_home
         self._payload_arm_away = payload_arm_away
         self._code = code
+        self._discovery_hash = discovery_hash
 
     @asyncio.coroutine
     def async_added_to_hass(self):
         """Subscribe mqtt events."""
-        yield from super().async_added_to_hass()
+        yield from MqttAvailability.async_added_to_hass(self)
+        yield from MqttDiscoveryUpdate.async_added_to_hass(self)
 
         @callback
         def message_received(topic, payload, qos):
diff --git a/tests/components/mqtt/test_discovery.py b/tests/components/mqtt/test_discovery.py
index 6de277eb48d..ee3460d46b5 100644
--- a/tests/components/mqtt/test_discovery.py
+++ b/tests/components/mqtt/test_discovery.py
@@ -209,3 +209,28 @@ def test_discovery_removal(hass, mqtt_mock, caplog):
 
     state = hass.states.get('switch.beer')
     assert state is None
+
+
+@asyncio.coroutine
+def test_discovery_removal_alarm(hass, mqtt_mock, caplog):
+    """Test removal of discovered alarm_control_panel."""
+    yield from async_start(hass, 'homeassistant', {})
+    data = (
+        '{ "name": "Beer",'
+        '  "status_topic": "test_topic",'
+        '  "command_topic": "test_topic" }'
+    )
+    async_fire_mqtt_message(hass,
+                            'homeassistant/alarm_control_panel/bla/config',
+                            data)
+    yield from hass.async_block_till_done()
+    state = hass.states.get('alarm_control_panel.beer')
+    assert state is not None
+    assert state.name == 'Beer'
+    async_fire_mqtt_message(hass,
+                            'homeassistant/alarm_control_panel/bla/config',
+                            '')
+    yield from hass.async_block_till_done()
+    yield from hass.async_block_till_done()
+    state = hass.states.get('alarm_control_panel.beer')
+    assert state is None