From 4a265f37e01682730905409754336adb9fc59b2e Mon Sep 17 00:00:00 2001
From: emontnemery <erik@montnemery.com>
Date: Tue, 25 Sep 2018 19:32:31 +0200
Subject: [PATCH] Remove discovered MQTT fan device when discovery topic is
 cleared (#16858)

---
 homeassistant/components/fan/mqtt.py | 24 ++++++++++++++++--------
 tests/components/fan/test_mqtt.py    | 25 ++++++++++++++++++++++++-
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/homeassistant/components/fan/mqtt.py b/homeassistant/components/fan/mqtt.py
index db3cfab3608..4ff8e1ec757 100644
--- a/homeassistant/components/fan/mqtt.py
+++ b/homeassistant/components/fan/mqtt.py
@@ -14,9 +14,9 @@ from homeassistant.const import (
     CONF_NAME, CONF_OPTIMISTIC, CONF_STATE, STATE_ON, STATE_OFF,
     CONF_PAYLOAD_OFF, CONF_PAYLOAD_ON)
 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
 from homeassistant.helpers.typing import HomeAssistantType, ConfigType
 from homeassistant.components.fan import (SPEED_LOW, SPEED_MEDIUM,
@@ -83,6 +83,10 @@ async def async_setup_platform(hass: HomeAssistantType, config: ConfigType,
     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([MqttFan(
         config.get(CONF_NAME),
         {
@@ -116,18 +120,20 @@ async def async_setup_platform(hass: HomeAssistantType, config: ConfigType,
         config.get(CONF_AVAILABILITY_TOPIC),
         config.get(CONF_PAYLOAD_AVAILABLE),
         config.get(CONF_PAYLOAD_NOT_AVAILABLE),
+        discovery_hash,
     )])
 
 
-class MqttFan(MqttAvailability, FanEntity):
+class MqttFan(MqttAvailability, MqttDiscoveryUpdate, FanEntity):
     """A MQTT fan component."""
 
     def __init__(self, name, topic, templates, qos, retain, payload,
                  speed_list, optimistic, availability_topic, payload_available,
-                 payload_not_available):
+                 payload_not_available, discovery_hash):
         """Initialize the MQTT fan."""
-        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._name = name
         self._topic = topic
         self._qos = qos
@@ -148,10 +154,12 @@ class MqttFan(MqttAvailability, FanEntity):
                                      is not None and SUPPORT_OSCILLATE)
         self._supported_features |= (topic[CONF_SPEED_STATE_TOPIC]
                                      is not None and SUPPORT_SET_SPEED)
+        self._discovery_hash = discovery_hash
 
     async def async_added_to_hass(self):
         """Subscribe to MQTT events."""
-        await super().async_added_to_hass()
+        await MqttAvailability.async_added_to_hass(self)
+        await MqttDiscoveryUpdate.async_added_to_hass(self)
 
         templates = {}
         for key, tpl in list(self._templates.items()):
diff --git a/tests/components/fan/test_mqtt.py b/tests/components/fan/test_mqtt.py
index 7f69e56218b..5f59ebe4a76 100644
--- a/tests/components/fan/test_mqtt.py
+++ b/tests/components/fan/test_mqtt.py
@@ -3,10 +3,12 @@ import unittest
 
 from homeassistant.setup import setup_component
 from homeassistant.components import fan
+from homeassistant.components.mqtt.discovery import async_start
 from homeassistant.const import ATTR_ASSUMED_STATE, STATE_UNAVAILABLE
 
 from tests.common import (
-    mock_mqtt_component, fire_mqtt_message, get_test_home_assistant)
+    mock_mqtt_component, async_fire_mqtt_message, fire_mqtt_message,
+    get_test_home_assistant)
 
 
 class TestMqttFan(unittest.TestCase):
@@ -102,3 +104,24 @@ class TestMqttFan(unittest.TestCase):
 
         state = self.hass.states.get('fan.test')
         self.assertNotEqual(STATE_UNAVAILABLE, state.state)
+
+
+async def test_discovery_removal_fan(hass, mqtt_mock, caplog):
+    """Test removal of discovered fan."""
+    await async_start(hass, 'homeassistant', {})
+    data = (
+        '{ "name": "Beer",'
+        '  "command_topic": "test_topic" }'
+    )
+    async_fire_mqtt_message(hass, 'homeassistant/fan/bla/config',
+                            data)
+    await hass.async_block_till_done()
+    state = hass.states.get('fan.beer')
+    assert state is not None
+    assert state.name == 'Beer'
+    async_fire_mqtt_message(hass, 'homeassistant/fan/bla/config',
+                            '')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    state = hass.states.get('fan.beer')
+    assert state is None