From 203701bc7c1069017759c6ff21107a1e96b08e95 Mon Sep 17 00:00:00 2001
From: emontnemery <erik@montnemery.com>
Date: Tue, 8 Jan 2019 16:51:03 +0100
Subject: [PATCH] Cleanup if discovered mqtt fan can't be added (#19741)

* Cleanup if discovered mqtt fan can't be added
---
 homeassistant/components/fan/mqtt.py | 19 ++++++++-----
 tests/components/fan/test_mqtt.py    | 40 +++++++++++++++++++++++++---
 2 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/homeassistant/components/fan/mqtt.py b/homeassistant/components/fan/mqtt.py
index 932a4584e2f..3adb821eba5 100644
--- a/homeassistant/components/fan/mqtt.py
+++ b/homeassistant/components/fan/mqtt.py
@@ -25,7 +25,8 @@ from homeassistant.components.fan import (SPEED_LOW, SPEED_MEDIUM,
                                           SPEED_HIGH, FanEntity,
                                           SUPPORT_SET_SPEED, SUPPORT_OSCILLATE,
                                           SPEED_OFF, ATTR_SPEED)
-from homeassistant.components.mqtt.discovery import MQTT_DISCOVERY_NEW
+from homeassistant.components.mqtt.discovery import (
+    MQTT_DISCOVERY_NEW, clear_discovery_hash)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -86,23 +87,29 @@ PLATFORM_SCHEMA = mqtt.MQTT_RW_PLATFORM_SCHEMA.extend({
 async def async_setup_platform(hass: HomeAssistantType, config: ConfigType,
                                async_add_entities, discovery_info=None):
     """Set up MQTT fan through configuration.yaml."""
-    await _async_setup_entity(hass, config, async_add_entities)
+    await _async_setup_entity(config, async_add_entities)
 
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
     """Set up MQTT fan dynamically through MQTT discovery."""
     async def async_discover(discovery_payload):
         """Discover and add a MQTT fan."""
-        config = PLATFORM_SCHEMA(discovery_payload)
-        await _async_setup_entity(hass, config, async_add_entities,
-                                  discovery_payload[ATTR_DISCOVERY_HASH])
+        try:
+            discovery_hash = discovery_payload[ATTR_DISCOVERY_HASH]
+            config = PLATFORM_SCHEMA(discovery_payload)
+            await _async_setup_entity(config, async_add_entities,
+                                      discovery_hash)
+        except Exception:
+            if discovery_hash:
+                clear_discovery_hash(hass, discovery_hash)
+            raise
 
     async_dispatcher_connect(
         hass, MQTT_DISCOVERY_NEW.format(fan.DOMAIN, 'mqtt'),
         async_discover)
 
 
-async def _async_setup_entity(hass, config, async_add_entities,
+async def _async_setup_entity(config, async_add_entities,
                               discovery_hash=None):
     """Set up the MQTT fan."""
     async_add_entities([MqttFan(
diff --git a/tests/components/fan/test_mqtt.py b/tests/components/fan/test_mqtt.py
index 405c20196a0..fea6f6dda74 100644
--- a/tests/components/fan/test_mqtt.py
+++ b/tests/components/fan/test_mqtt.py
@@ -3,7 +3,7 @@ import json
 from unittest.mock import ANY
 
 from homeassistant.setup import async_setup_component
-from homeassistant.components import fan
+from homeassistant.components import fan, mqtt
 from homeassistant.components.mqtt.discovery import async_start
 from homeassistant.const import ATTR_ASSUMED_STATE, STATE_UNAVAILABLE
 
@@ -111,7 +111,7 @@ async def test_custom_availability_payload(hass, mqtt_mock):
 
 async def test_discovery_removal_fan(hass, mqtt_mock, caplog):
     """Test removal of discovered fan."""
-    entry = MockConfigEntry(domain='mqtt')
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
     await async_start(hass, 'homeassistant', {}, entry)
     data = (
         '{ "name": "Beer",'
@@ -132,8 +132,8 @@ async def test_discovery_removal_fan(hass, mqtt_mock, caplog):
 
 
 async def test_discovery_update_fan(hass, mqtt_mock, caplog):
-    """Test removal of discovered fan."""
-    entry = MockConfigEntry(domain='mqtt')
+    """Test update of discovered fan."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
     await async_start(hass, 'homeassistant', {}, entry)
     data1 = (
         '{ "name": "Beer",'
@@ -163,6 +163,38 @@ async def test_discovery_update_fan(hass, mqtt_mock, caplog):
     assert state is None
 
 
+async def test_discovery_broken(hass, mqtt_mock, caplog):
+    """Test handling of bad discovery message."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    await async_start(hass, 'homeassistant', {}, entry)
+
+    data1 = (
+        '{ "name": "Beer" }'
+    )
+    data2 = (
+        '{ "name": "Milk",'
+        '  "command_topic": "test_topic" }'
+    )
+
+    async_fire_mqtt_message(hass, 'homeassistant/fan/bla/config',
+                            data1)
+    await hass.async_block_till_done()
+
+    state = hass.states.get('fan.beer')
+    assert state is None
+
+    async_fire_mqtt_message(hass, 'homeassistant/fan/bla/config',
+                            data2)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get('fan.milk')
+    assert state is not None
+    assert state.name == 'Milk'
+    state = hass.states.get('fan.beer')
+    assert state is None
+
+
 async def test_unique_id(hass):
     """Test unique_id option only creates one fan per id."""
     await async_mock_mqtt_component(hass)