From 9290f245bfd51d1c3b9bd81053c0116a49c90d25 Mon Sep 17 00:00:00 2001
From: Otto Winter <otto@otto-winter.com>
Date: Mon, 8 Oct 2018 14:59:04 +0200
Subject: [PATCH] Convert MQTT fan to config entry (#17247)

---
 homeassistant/components/fan/mqtt.py       | 28 ++++++++++++++++------
 homeassistant/components/mqtt/discovery.py |  1 +
 tests/components/fan/test_mqtt.py          |  5 ++--
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/homeassistant/components/fan/mqtt.py b/homeassistant/components/fan/mqtt.py
index 3e1ad2704e7..bc45754f698 100644
--- a/homeassistant/components/fan/mqtt.py
+++ b/homeassistant/components/fan/mqtt.py
@@ -10,7 +10,7 @@ from typing import Optional
 import voluptuous as vol
 
 from homeassistant.core import callback
-from homeassistant.components import mqtt
+from homeassistant.components import fan, mqtt
 from homeassistant.const import (
     CONF_NAME, CONF_OPTIMISTIC, CONF_STATE, STATE_ON, STATE_OFF,
     CONF_PAYLOAD_OFF, CONF_PAYLOAD_ON)
@@ -19,11 +19,13 @@ from homeassistant.components.mqtt import (
     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.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.typing import HomeAssistantType, ConfigType
 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
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -82,14 +84,26 @@ 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 the MQTT fan platform."""
-    if discovery_info is not None:
-        config = PLATFORM_SCHEMA(discovery_info)
+    """Set up MQTT fan through configuration.yaml."""
+    await _async_setup_entity(hass, config, async_add_entities)
 
-    discovery_hash = None
-    if discovery_info is not None and ATTR_DISCOVERY_HASH in discovery_info:
-        discovery_hash = discovery_info[ATTR_DISCOVERY_HASH]
 
+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])
+
+    async_dispatcher_connect(
+        hass, MQTT_DISCOVERY_NEW.format(fan.DOMAIN, 'mqtt'),
+        async_discover)
+
+
+async def _async_setup_entity(hass, config, async_add_entities,
+                              discovery_hash=None):
+    """Set up the MQTT fan."""
     async_add_entities([MqttFan(
         config.get(CONF_NAME),
         {
diff --git a/homeassistant/components/mqtt/discovery.py b/homeassistant/components/mqtt/discovery.py
index a762978a330..62f500d9952 100644
--- a/homeassistant/components/mqtt/discovery.py
+++ b/homeassistant/components/mqtt/discovery.py
@@ -49,6 +49,7 @@ CONFIG_ENTRY_PLATFORMS = {
     'switch': ['mqtt'],
     'climate': ['mqtt'],
     'alarm_control_panel': ['mqtt'],
+    'fan': ['mqtt'],
 }
 
 ALREADY_DISCOVERED = 'mqtt_discovered_components'
diff --git a/tests/components/fan/test_mqtt.py b/tests/components/fan/test_mqtt.py
index 7434e5aa1c9..feb2ff6904d 100644
--- a/tests/components/fan/test_mqtt.py
+++ b/tests/components/fan/test_mqtt.py
@@ -8,7 +8,7 @@ from homeassistant.const import ATTR_ASSUMED_STATE, STATE_UNAVAILABLE
 
 from tests.common import (
     mock_mqtt_component, async_fire_mqtt_message, fire_mqtt_message,
-    get_test_home_assistant, async_mock_mqtt_component)
+    get_test_home_assistant, async_mock_mqtt_component, MockConfigEntry)
 
 
 class TestMqttFan(unittest.TestCase):
@@ -108,7 +108,8 @@ class TestMqttFan(unittest.TestCase):
 
 async def test_discovery_removal_fan(hass, mqtt_mock, caplog):
     """Test removal of discovered fan."""
-    await async_start(hass, 'homeassistant', {})
+    entry = MockConfigEntry(domain='mqtt')
+    await async_start(hass, 'homeassistant', {}, entry)
     data = (
         '{ "name": "Beer",'
         '  "command_topic": "test_topic" }'