From 8de863ecf1214f2fbf30e80b975d5efe7ea13160 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Tue, 26 May 2020 17:12:22 +0200
Subject: [PATCH] Let PAHO MQTT client handle connection to MQTT server
 (#35983)

* Let PAHO client handle connection to MQTT server
---
 homeassistant/components/mqtt/__init__.py | 28 ++++++++---------------
 tests/components/mqtt/test_init.py        | 12 +++++-----
 2 files changed, 15 insertions(+), 25 deletions(-)

diff --git a/homeassistant/components/mqtt/__init__.py b/homeassistant/components/mqtt/__init__.py
index 25b2b0381ea..64b25ad9486 100644
--- a/homeassistant/components/mqtt/__init__.py
+++ b/homeassistant/components/mqtt/__init__.py
@@ -29,11 +29,7 @@ from homeassistant.const import (
     EVENT_HOMEASSISTANT_STOP,
 )
 from homeassistant.core import Event, ServiceCall, callback
-from homeassistant.exceptions import (
-    ConfigEntryNotReady,
-    HomeAssistantError,
-    Unauthorized,
-)
+from homeassistant.exceptions import HomeAssistantError, Unauthorized
 from homeassistant.helpers import config_validation as cv, event, template
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import Entity
@@ -649,13 +645,7 @@ async def async_setup_entry(hass, entry):
         tls_version=tls_version,
     )
 
-    result: str = await hass.data[DATA_MQTT].async_connect()
-
-    if result == CONNECTION_FAILED:
-        return False
-
-    if result == CONNECTION_FAILED_RECOVERABLE:
-        raise ConfigEntryNotReady
+    await hass.data[DATA_MQTT].async_connect()
 
     async def async_stop_mqtt(_event: Event):
         """Stop MQTT component."""
@@ -835,15 +825,14 @@ class MQTT:
                 self._mqttc.connect, self.broker, self.port, self.keepalive
             )
         except OSError as err:
-            _LOGGER.error("Failed to connect due to exception: %s", err)
-            return CONNECTION_FAILED_RECOVERABLE
+            _LOGGER.error("Failed to connect to MQTT server due to exception: %s", err)
 
-        if result != 0:
-            _LOGGER.error("Failed to connect: %s", mqtt.error_string(result))
-            return CONNECTION_FAILED
+        if result is not None and result != 0:
+            _LOGGER.error(
+                "Failed to connect to MQTT server: %s", mqtt.error_string(result)
+            )
 
         self._mqttc.loop_start()
-        return CONNECTION_SUCCESS
 
     async def async_disconnect(self):
         """Stop the MQTT client."""
@@ -933,6 +922,7 @@ class MQTT:
             return
 
         self.connected = True
+        _LOGGER.info("Connected to MQTT server (%s)", result_code)
 
         # Group subscriptions to only re-subscribe once for each topic.
         keyfunc = attrgetter("topic")
@@ -999,7 +989,7 @@ class MQTT:
     def _mqtt_on_disconnect(self, _mqttc, _userdata, result_code: int) -> None:
         """Disconnected callback."""
         self.connected = False
-        _LOGGER.warning("Disconnected from MQTT (%s).", result_code)
+        _LOGGER.warning("Disconnected from MQTT server (%s)", result_code)
 
 
 def _raise_on_error(result_code: int) -> None:
diff --git a/tests/components/mqtt/test_init.py b/tests/components/mqtt/test_init.py
index 9ec5e09f276..3626c5a746c 100644
--- a/tests/components/mqtt/test_init.py
+++ b/tests/components/mqtt/test_init.py
@@ -18,7 +18,6 @@ from homeassistant.const import (
     TEMP_CELSIUS,
 )
 from homeassistant.core import callback
-from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import device_registry
 from homeassistant.setup import async_setup_component
 from homeassistant.util.dt import utcnow
@@ -678,23 +677,24 @@ async def test_setup_embedded_with_embedded(hass):
         assert _start.call_count == 1
 
 
-async def test_setup_fails_if_no_connect_broker(hass):
+async def test_setup_logs_error_if_no_connect_broker(hass, caplog):
     """Test for setup failure if connection to broker is missing."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN, data={mqtt.CONF_BROKER: "test-broker"})
 
     with patch("paho.mqtt.client.Client") as mock_client:
         mock_client().connect = lambda *args: 1
-        assert not await mqtt.async_setup_entry(hass, entry)
+        assert await mqtt.async_setup_entry(hass, entry)
+        assert "Failed to connect to MQTT server:" in caplog.text
 
 
-async def test_setup_raises_ConfigEntryNotReady_if_no_connect_broker(hass):
+async def test_setup_raises_ConfigEntryNotReady_if_no_connect_broker(hass, caplog):
     """Test for setup failure if connection to broker is missing."""
     entry = MockConfigEntry(domain=mqtt.DOMAIN, data={mqtt.CONF_BROKER: "test-broker"})
 
     with patch("paho.mqtt.client.Client") as mock_client:
         mock_client().connect = MagicMock(side_effect=OSError("Connection error"))
-        with pytest.raises(ConfigEntryNotReady):
-            await mqtt.async_setup_entry(hass, entry)
+        assert await mqtt.async_setup_entry(hass, entry)
+        assert "Failed to connect to MQTT server due to exception:" in caplog.text
 
 
 async def test_setup_uses_certificate_on_certificate_set_to_auto(hass, mock_mqtt):