From 0dbfd77402777f195fbc87364e185d7914dd3f4b Mon Sep 17 00:00:00 2001
From: emontnemery <erik@montnemery.com>
Date: Tue, 25 Sep 2018 17:15:39 +0200
Subject: [PATCH] Remove discovered MQTT climate device when discovery topic is
 cleared (#16856)

---
 homeassistant/components/climate/mqtt.py | 27 ++++++++++++++++--------
 tests/components/climate/test_mqtt.py    | 24 ++++++++++++++++++++-
 2 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/homeassistant/components/climate/mqtt.py b/homeassistant/components/climate/mqtt.py
index 9e227e002b5..66f76ac1aaa 100644
--- a/homeassistant/components/climate/mqtt.py
+++ b/homeassistant/components/climate/mqtt.py
@@ -21,8 +21,9 @@ from homeassistant.components.climate import (
 from homeassistant.const import (
     STATE_ON, STATE_OFF, ATTR_TEMPERATURE, CONF_NAME, CONF_VALUE_TEMPLATE)
 from homeassistant.components.mqtt import (
-    CONF_AVAILABILITY_TOPIC, CONF_QOS, CONF_RETAIN, CONF_PAYLOAD_AVAILABLE,
-    CONF_PAYLOAD_NOT_AVAILABLE, MQTT_BASE_PLATFORM_SCHEMA, MqttAvailability)
+    ATTR_DISCOVERY_HASH, CONF_AVAILABILITY_TOPIC, CONF_QOS, CONF_RETAIN,
+    CONF_PAYLOAD_AVAILABLE, CONF_PAYLOAD_NOT_AVAILABLE,
+    MQTT_BASE_PLATFORM_SCHEMA, MqttAvailability, MqttDiscoveryUpdate)
 import homeassistant.helpers.config_validation as cv
 from homeassistant.components.fan import (SPEED_LOW, SPEED_MEDIUM,
                                           SPEED_HIGH)
@@ -153,6 +154,10 @@ def async_setup_platform(hass, config, async_add_entities,
         value_templates[key] = config.get(key)
         value_templates[key].hass = hass
 
+    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([
         MqttClimate(
             hass,
@@ -194,11 +199,12 @@ def async_setup_platform(hass, config, async_add_entities,
             config.get(CONF_PAYLOAD_AVAILABLE),
             config.get(CONF_PAYLOAD_NOT_AVAILABLE),
             config.get(CONF_MIN_TEMP),
-            config.get(CONF_MAX_TEMP))
-    ])
+            config.get(CONF_MAX_TEMP),
+            discovery_hash,
+        )])
 
 
-class MqttClimate(MqttAvailability, ClimateDevice):
+class MqttClimate(MqttAvailability, MqttDiscoveryUpdate, ClimateDevice):
     """Representation of an MQTT climate device."""
 
     def __init__(self, hass, name, topic, value_templates, qos, retain,
@@ -207,10 +213,11 @@ class MqttClimate(MqttAvailability, ClimateDevice):
                  current_swing_mode, current_operation, aux, send_if_off,
                  payload_on, payload_off, availability_topic,
                  payload_available, payload_not_available,
-                 min_temp, max_temp):
+                 min_temp, max_temp, discovery_hash):
         """Initialize the climate device."""
-        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.hass = hass
         self._name = name
         self._topic = topic
@@ -235,11 +242,13 @@ class MqttClimate(MqttAvailability, ClimateDevice):
         self._payload_off = payload_off
         self._min_temp = min_temp
         self._max_temp = max_temp
+        self._discovery_hash = discovery_hash
 
     @asyncio.coroutine
     def async_added_to_hass(self):
         """Handle being added to home assistant."""
-        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 handle_current_temp_received(topic, payload, qos):
diff --git a/tests/components/climate/test_mqtt.py b/tests/components/climate/test_mqtt.py
index f46a23e4f97..6c43c4d9dbe 100644
--- a/tests/components/climate/test_mqtt.py
+++ b/tests/components/climate/test_mqtt.py
@@ -12,8 +12,10 @@ from homeassistant.components.climate import (
     SUPPORT_OPERATION_MODE, SUPPORT_TARGET_TEMPERATURE,
     SUPPORT_FAN_MODE, SUPPORT_SWING_MODE, SUPPORT_HOLD_MODE,
     SUPPORT_AWAY_MODE, SUPPORT_AUX_HEAT, DEFAULT_MIN_TEMP, DEFAULT_MAX_TEMP)
+from homeassistant.components.mqtt.discovery import async_start
 from tests.common import (get_test_home_assistant, mock_mqtt_component,
-                          fire_mqtt_message, mock_component)
+                          async_fire_mqtt_message, fire_mqtt_message,
+                          mock_component)
 
 ENTITY_CLIMATE = 'climate.test'
 
@@ -649,3 +651,23 @@ class TestMQTTClimate(unittest.TestCase):
 
         self.assertIsInstance(max_temp, float)
         self.assertEqual(60, max_temp)
+
+
+async def test_discovery_removal_climate(hass, mqtt_mock, caplog):
+    """Test removal of discovered climate."""
+    await async_start(hass, 'homeassistant', {})
+    data = (
+        '{ "name": "Beer" }'
+    )
+    async_fire_mqtt_message(hass, 'homeassistant/climate/bla/config',
+                            data)
+    await hass.async_block_till_done()
+    state = hass.states.get('climate.beer')
+    assert state is not None
+    assert state.name == 'Beer'
+    async_fire_mqtt_message(hass, 'homeassistant/climate/bla/config',
+                            '')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    state = hass.states.get('climate.beer')
+    assert state is None