Allow MQTT discovery (#16842)

homeassistant/components/mqtt/.translations/en.json

@@ -10,6 +10,7 @@
             "broker": {
                 "data": {
                     "broker": "Broker",
+                    "discovery": "Enable discovery",
                     "password": "Password",
                     "port": "Port",
                     "username": "Username"

homeassistant/components/mqtt/__init__.py

@@ -36,7 +36,7 @@ from homeassistant.util.async_ import (
 
 # Loading the config flow file will register the flow
 from . import config_flow  # noqa  # pylint: disable=unused-import
-from .const import CONF_BROKER
+from .const import CONF_BROKER, CONF_DISCOVERY, DEFAULT_DISCOVERY
 from .server import HBMQTT_CONFIG_SCHEMA
 
 REQUIREMENTS = ['paho-mqtt==1.4.0']
@@ -47,13 +47,13 @@ DOMAIN = 'mqtt'
 
 DATA_MQTT = 'mqtt'
 DATA_MQTT_CONFIG = 'mqtt_config'
+DATA_MQTT_HASS_CONFIG = 'mqtt_hass_config'
 
 SERVICE_PUBLISH = 'publish'
 
 CONF_EMBEDDED = 'embedded'
 
 CONF_CLIENT_ID = 'client_id'
-CONF_DISCOVERY = 'discovery'
 CONF_DISCOVERY_PREFIX = 'discovery_prefix'
 CONF_KEEPALIVE = 'keepalive'
 CONF_CERTIFICATE = 'certificate'
@@ -81,7 +81,6 @@ DEFAULT_KEEPALIVE = 60
 DEFAULT_QOS = 0
 DEFAULT_RETAIN = False
 DEFAULT_PROTOCOL = PROTOCOL_311
-DEFAULT_DISCOVERY = False
 DEFAULT_DISCOVERY_PREFIX = 'homeassistant'
 DEFAULT_TLS_PROTOCOL = 'auto'
 DEFAULT_PAYLOAD_AVAILABLE = 'online'
@@ -321,23 +320,21 @@ async def _async_setup_server(hass: HomeAssistantType, config: ConfigType):
     return broker_config
 
 
-async def _async_setup_discovery(hass: HomeAssistantType,
+async def _async_setup_discovery(hass: HomeAssistantType, conf: ConfigType,
-                                 config: ConfigType) -> bool:
+                                 hass_config: ConfigType) -> bool:
     """Try to start the discovery of MQTT devices.
 
     This method is a coroutine.
     """
-    conf = config.get(DOMAIN, {})  # type: ConfigType
-
     discovery = await async_prepare_setup_platform(
-        hass, config, DOMAIN, 'discovery')
+        hass, hass_config, DOMAIN, 'discovery')
 
     if discovery is None:
         _LOGGER.error("Unable to load MQTT discovery")
         return False
 
     success = await discovery.async_start(
-        hass, conf[CONF_DISCOVERY_PREFIX], config)  # type: bool
+        hass, conf[CONF_DISCOVERY_PREFIX], hass_config)  # type: bool
 
     return success
 
@@ -346,6 +343,11 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
     """Start the MQTT protocol service."""
     conf = config.get(DOMAIN)  # type: Optional[ConfigType]
 
+    # We need this because discovery can cause components to be set up and
+    # otherwise it will not load the users config.
+    # This needs a better solution.
+    hass.data[DATA_MQTT_HASS_CONFIG] = config
+
     if conf is None:
         # If we have a config entry, setup is done by that config entry.
         # If there is no config entry, this should fail.
@@ -390,13 +392,6 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType) -> bool:
             data={}
         ))
 
-    if conf.get(CONF_DISCOVERY):
-        async def async_setup_discovery(event):
-            await _async_setup_discovery(hass, config)
-
-        hass.bus.async_listen_once(
-            EVENT_HOMEASSISTANT_START, async_setup_discovery)
-
     return True
 
 
@@ -528,6 +523,14 @@ async def async_setup_entry(hass, entry):
         DOMAIN, SERVICE_PUBLISH, async_publish_service,
         schema=MQTT_PUBLISH_SCHEMA)
 
+    if conf.get(CONF_DISCOVERY):
+        async def async_setup_discovery(event):
+            await _async_setup_discovery(
+                hass, conf, hass.data[DATA_MQTT_HASS_CONFIG])
+
+        hass.bus.async_listen_once(
+            EVENT_HOMEASSISTANT_START, async_setup_discovery)
+
     return True
 
 

homeassistant/components/mqtt/config_flow.py

@@ -7,7 +7,7 @@ import voluptuous as vol
 from homeassistant import config_entries
 from homeassistant.const import CONF_PASSWORD, CONF_PORT, CONF_USERNAME
 
-from .const import CONF_BROKER
+from .const import CONF_BROKER, CONF_DISCOVERY, DEFAULT_DISCOVERY
 
 
 @config_entries.HANDLERS.register('mqtt')
@@ -44,6 +44,7 @@ class FlowHandler(config_entries.ConfigFlow):
         fields[vol.Required(CONF_PORT, default=1883)] = vol.Coerce(int)
         fields[vol.Optional(CONF_USERNAME)] = str
         fields[vol.Optional(CONF_PASSWORD)] = str
+        fields[vol.Optional(CONF_DISCOVERY, default=DEFAULT_DISCOVERY)] = bool
 
         return self.async_show_form(
             step_id='broker', data_schema=vol.Schema(fields), errors=errors)

homeassistant/components/mqtt/const.py

@@ -1,2 +1,4 @@
 """Constants used by multiple MQTT modules."""
 CONF_BROKER = 'broker'
+CONF_DISCOVERY = 'discovery'
+DEFAULT_DISCOVERY = False

homeassistant/components/mqtt/strings.json

@@ -9,7 +9,8 @@
           "broker": "Broker",
           "port": "Port",
           "username": "Username",
-          "password": "Password"
+          "password": "Password",
+          "discovery": "Enable discovery"
         }
       }
     },

tests/components/mqtt/test_config_flow.py

@@ -41,6 +41,11 @@ async def test_user_connection_works(hass, mock_try_connection,
     )
 
     assert result['type'] == 'create_entry'
+    assert result['result'].data == {
+        'broker': '127.0.0.1',
+        'port': 1883,
+        'discovery': False,
+    }
     # Check we tried the connection
     assert len(mock_try_connection.mock_calls) == 1
     # Check config entry got setup