Raise ConfigEntryNotReady for MQTT connection exception (#22540)

* Raise ConfigEntryNotReady for connection exception

Raise ConfigEntryNotReady for the connection exception like if the MQTT Server container/device is being restarted or was unavailable on boot.

* Add new exception

* grammar fix

* Possibly resolve hound comments

* raise `ConfigEntryNotReady` for mqtt connection error

* revert exceptions.py

* Update exceptions.py

* modify test to handle exception

* use constants to control exception scope

* Raise ConfigEntryNotReady for connection exception

Raise ConfigEntryNotReady for the connection exception like if the MQTT Server container/device is being restarted or was unavailable on boot.

* Add new exception

* Add new exception

* grammar fix

* Possibly resolve hound comments

* raise `ConfigEntryNotReady` for mqtt connection error

* revert exceptions.py

* Update exceptions.py

* modify test to handle exception

* use constants to control exception scope

* revert test change as it's not the same thing

* Update test_init.py

* Add test for MQTT OSError

* revert file changes from a bad rebase

* Rewrite test with valid syntax

* rewrite test to be less ambiguous

* add empty line

* add back 'axis'

* Remove empty line

* Update tests and undo merge from earlier

* correctly restore test for no connect broker

* fix test mock correctly

* line was too long. hit enter.

homeassistant/components/mqtt/__init__.py

@@ -23,7 +23,8 @@ from homeassistant.const import (
     CONF_PROTOCOL, CONF_USERNAME, CONF_VALUE_TEMPLATE,
     EVENT_HOMEASSISTANT_STOP)
 from homeassistant.core import Event, ServiceCall, callback
-from homeassistant.exceptions import HomeAssistantError, Unauthorized
+from homeassistant.exceptions import (
+    HomeAssistantError, Unauthorized, ConfigEntryNotReady)
 from homeassistant.helpers import config_validation as cv, template
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.typing import (
@@ -104,6 +105,10 @@ ATTR_DISCOVERY_HASH = 'discovery_hash'
 
 MAX_RECONNECT_WAIT = 300  # seconds
 
+CONNECTION_SUCCESS = 'connection_success'
+CONNECTION_FAILED = 'connection_failed'
+CONNECTION_FAILED_RECOVERABLE = 'connection_failed_recoverable'
+
 
 def valid_topic(value: Any) -> str:
     """Validate that this is a valid topic name/filter."""
@@ -569,11 +574,14 @@ async def async_setup_entry(hass, entry):
         tls_version=tls_version,
     )
 
-    success = await hass.data[DATA_MQTT].async_connect()  # type: bool
+    result = await hass.data[DATA_MQTT].async_connect()  # type: str
 
-    if not success:
+    if result == CONNECTION_FAILED:
         return False
 
+    if result == CONNECTION_FAILED_RECOVERABLE:
+        raise ConfigEntryNotReady
+
     async def async_stop_mqtt(event: Event):
         """Stop MQTT component."""
         await hass.data[DATA_MQTT].async_disconnect()
@@ -685,7 +693,7 @@ class MQTT:
             await self.hass.async_add_job(
                 self._mqttc.publish, topic, payload, qos, retain)
 
-    async def async_connect(self) -> bool:
+    async def async_connect(self) -> str:
         """Connect to the host. Does process messages yet.
 
         This method is a coroutine.
@@ -696,15 +704,15 @@ 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 False
+            return CONNECTION_FAILED_RECOVERABLE
 
         if result != 0:
             import paho.mqtt.client as mqtt
             _LOGGER.error("Failed to connect: %s", mqtt.error_string(result))
-            return False
+            return CONNECTION_FAILED
 
         self._mqttc.loop_start()
-        return True
+        return CONNECTION_SUCCESS
 
     @callback
     def async_disconnect(self):

tests/components/mqtt/test_init.py

@@ -12,6 +12,7 @@ from homeassistant.const import (
     ATTR_DOMAIN, ATTR_SERVICE, EVENT_CALL_SERVICE, EVENT_HOMEASSISTANT_STOP)
 from homeassistant.core import callback
 from homeassistant.setup import async_setup_component
+from homeassistant.exceptions import ConfigEntryNotReady
 
 from tests.common import (
     MockConfigEntry, async_fire_mqtt_message, async_mock_mqtt_component,
@@ -621,6 +622,19 @@ async def test_setup_fails_if_no_connect_broker(hass):
         assert not await mqtt.async_setup_entry(hass, entry)
 
 
+async def test_setup_raises_ConfigEntryNotReady_if_no_connect_broker(hass):
+    """Test for setup failure if connection to broker is missing."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN, data={
+        mqtt.CONF_BROKER: 'test-broker'
+    })
+
+    with mock.patch('paho.mqtt.client.Client') as mock_client:
+        mock_client().connect = mock.Mock(
+            side_effect=OSError("Connection error"))
+        with pytest.raises(ConfigEntryNotReady):
+            await mqtt.async_setup_entry(hass, entry)
+
+
 async def test_setup_uses_certificate_on_certificate_set_to_auto(
         hass, mock_MQTT):
     """Test setup uses bundled certs when certificate is set to auto."""