Add Hass.io discovery to MQTT (#16962)

* Add Hass.io discovery to MQTT

* Update en.json

* Update __init__.py

* Update config_flow.py

* Update strings.json

* Update test_config_flow.py

* Lint

* Ensure we don't send bad data in config entry

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

@@ -17,8 +17,15 @@
                 },
                 "description": "Please enter the connection information of your MQTT broker.",
                 "title": "MQTT"
+            },
+            "hassio_confirm": {
+                "data": {
+                    "discovery": "Enable discovery"
+                },
+                "description": "Do you want to configure Home Assistant to connect to the MQTT broker provided by the hass.io add-on {addon}?",
+                "title": "MQTT Broker via Hass.io add-on"
             }
         },
         "title": "MQTT"
     }
-}
+}

homeassistant/components/mqtt/__init__.py

@@ -411,7 +411,7 @@ async def async_setup_entry(hass, entry):
     # If user didn't have configuration.yaml config, generate defaults
     if conf is None:
         conf = CONFIG_SCHEMA({
-            DOMAIN: entry.data
+            DOMAIN: entry.data,
         })[DOMAIN]
     elif any(key in conf for key in entry.data):
         _LOGGER.warning(

homeassistant/components/mqtt/config_flow.py

@@ -5,7 +5,8 @@ import queue
 import voluptuous as vol
 
 from homeassistant import config_entries
-from homeassistant.const import CONF_PASSWORD, CONF_PORT, CONF_USERNAME
+from homeassistant.const import (
+    CONF_PASSWORD, CONF_PORT, CONF_USERNAME, CONF_PROTOCOL)
 
 from .const import CONF_BROKER, CONF_DISCOVERY, DEFAULT_DISCOVERY
 
@@ -17,6 +18,8 @@ class FlowHandler(config_entries.ConfigFlow):
     VERSION = 1
     CONNECTION_CLASS = config_entries.CONN_CLASS_LOCAL_PUSH
 
+    _hassio_discovery = None
+
     async def async_step_user(self, user_input=None):
         """Handle a flow initialized by the user."""
         if self._async_current_entries():
@@ -60,11 +63,65 @@ class FlowHandler(config_entries.ConfigFlow):
 
         return self.async_create_entry(title='configuration.yaml', data={})
 
+    async def async_step_hassio(self, user_input=None):
+        """Receive a Hass.io discovery."""
+        if self._async_current_entries():
+            return self.async_abort(reason='single_instance_allowed')
 
-def try_connection(broker, port, username, password):
+        self._hassio_discovery = user_input
+
+        return await self.async_step_hassio_confirm()
+
+    async def async_step_hassio_confirm(self, user_input=None):
+        """Confirm a Hass.io discovery."""
+        errors = {}
+
+        if user_input is not None:
+            data = self._hassio_discovery
+            can_connect = await self.hass.async_add_executor_job(
+                try_connection,
+                data[CONF_BROKER],
+                data[CONF_PORT],
+                data.get(CONF_USERNAME),
+                data.get(CONF_PASSWORD),
+                data.get(CONF_PROTOCOL)
+            )
+
+            if can_connect:
+                return self.async_create_entry(
+                    title=data['addon'], data={
+                        CONF_BROKER: data[CONF_BROKER],
+                        CONF_PORT: data[CONF_PORT],
+                        CONF_USERNAME: data.get(CONF_USERNAME),
+                        CONF_PASSWORD: data.get(CONF_PASSWORD),
+                        CONF_PROTOCOL: data.get(CONF_PROTOCOL),
+                        CONF_DISCOVERY: user_input[CONF_DISCOVERY],
+                    })
+
+            errors['base'] = 'cannot_connect'
+
+        return self.async_show_form(
+            step_id='hassio_confirm',
+            description_placeholders={
+                'addon': self._hassio_discovery['addon']
+            },
+            data_schema=vol.Schema({
+                vol.Optional(CONF_DISCOVERY, default=DEFAULT_DISCOVERY): bool
+            }),
+            errors=errors,
+        )
+
+
+def try_connection(broker, port, username, password, protocol='3.1'):
     """Test if we can connect to an MQTT broker."""
     import paho.mqtt.client as mqtt
-    client = mqtt.Client()
+
+    if protocol == '3.1':
+        proto = mqtt.MQTTv31
+    else:
+        proto = mqtt.MQTTv311
+
+    client = mqtt.Client(protocol=proto)
     if username and password:
         client.username_pw_set(username, password)
 

homeassistant/components/mqtt/strings.json

@@ -12,6 +12,13 @@
           "password": "Password",
           "discovery": "Enable discovery"
         }
+      },
+      "hassio_confirm": {
+        "title": "MQTT Broker via Hass.io add-on",
+        "description": "Do you want to configure Home Assistant to connect to the MQTT broker provided by the hass.io add-on {addon}?",
+        "data": {
+          "discovery": "Enable discovery"
+        }
       }
     },
     "abort": {

tests/components/mqtt/test_config_flow.py

@@ -5,7 +5,7 @@ import pytest
 
 from homeassistant.setup import async_setup_component
 
-from tests.common import mock_coro
+from tests.common import mock_coro, MockConfigEntry
 
 
 @pytest.fixture(autouse=True)
@@ -88,3 +88,67 @@ async def test_manual_config_set(hass, mock_try_connection,
     result = await hass.config_entries.flow.async_init(
         'mqtt', context={'source': 'user'})
     assert result['type'] == 'abort'
+
+
+async def test_user_single_instance(hass):
+    """Test we only allow a single config flow."""
+    MockConfigEntry(domain='mqtt').add_to_hass(hass)
+
+    result = await hass.config_entries.flow.async_init(
+        'mqtt', context={'source': 'user'})
+    assert result['type'] == 'abort'
+    assert result['reason'] == 'single_instance_allowed'
+
+
+async def test_hassio_single_instance(hass):
+    """Test we only allow a single config flow."""
+    MockConfigEntry(domain='mqtt').add_to_hass(hass)
+
+    result = await hass.config_entries.flow.async_init(
+        'mqtt', context={'source': 'hassio'})
+    assert result['type'] == 'abort'
+    assert result['reason'] == 'single_instance_allowed'
+
+
+async def test_hassio_confirm(hass, mock_try_connection,
+                              mock_finish_setup):
+    """Test we can finish a config flow."""
+    mock_try_connection.return_value = True
+
+    result = await hass.config_entries.flow.async_init(
+        'mqtt',
+        data={
+            'addon': 'Mock Addon',
+            'broker': 'mock-broker',
+            'port': 1883,
+            'username': 'mock-user',
+            'password': 'mock-pass',
+            'protocol': '3.1.1'
+        },
+        context={'source': 'hassio'}
+    )
+    assert result['type'] == 'form'
+    assert result['step_id'] == 'hassio_confirm'
+    assert result['description_placeholders'] == {
+        'addon': 'Mock Addon',
+    }
+
+    result = await hass.config_entries.flow.async_configure(
+        result['flow_id'], {
+            'discovery': True,
+        }
+    )
+
+    assert result['type'] == 'create_entry'
+    assert result['result'].data == {
+        'broker': 'mock-broker',
+        'port': 1883,
+        'username': 'mock-user',
+        'password': 'mock-pass',
+        'protocol': '3.1.1',
+        'discovery': True,
+    }
+    # Check we tried the connection
+    assert len(mock_try_connection.mock_calls) == 1
+    # Check config entry got setup
+    assert len(mock_finish_setup.mock_calls) == 1