Fix some ESPHome race conditions (#19772)

* Fix some ESPHome race conditions

* Remove debug

* Update requirements_all.txt

* 🚑 Fix IDE line length settings

homeassistant/components/esphome/__init__.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, List, Optional, TYPE_CHECKING, Callable
 import attr
 import voluptuous as vol
 
+from homeassistant import const
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT, \
     EVENT_HOMEASSISTANT_STOP
@@ -30,7 +31,7 @@ if TYPE_CHECKING:
         ServiceCall
 
 DOMAIN = 'esphome'
-REQUIREMENTS = ['aioesphomeapi==1.3.0']
+REQUIREMENTS = ['aioesphomeapi==1.4.0']
 
 
 DISPATCHER_UPDATE_ENTITY = 'esphome_{entry_id}_update_{component_key}_{key}'
@@ -161,8 +162,8 @@ async def async_setup_entry(hass: HomeAssistantType,
     port = entry.data[CONF_PORT]
     password = entry.data[CONF_PASSWORD]
 
-    cli = APIClient(hass.loop, host, port, password)
+    cli = APIClient(hass.loop, host, port, password,
-    await cli.start()
+                    client_info="Home Assistant {}".format(const.__version__))
 
     # Store client in per-config-entry hass.data
     store = Store(hass, STORAGE_VERSION, STORAGE_KEY.format(entry.entry_id),
@@ -181,8 +182,6 @@ async def async_setup_entry(hass: HomeAssistantType,
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, on_stop)
     )
 
-    try_connect = await _setup_auto_reconnect_logic(hass, cli, entry, host)
-
     @callback
     def async_on_state(state: 'EntityState') -> None:
         """Send dispatcher updates when a new state is received."""
@@ -247,7 +246,8 @@ async def async_setup_entry(hass: HomeAssistantType,
             # Re-connection logic will trigger after this
             await cli.disconnect()
 
-    cli.on_login = on_login
+    try_connect = await _setup_auto_reconnect_logic(hass, cli, entry, host,
+                                                    on_login)
 
     # This is a bit of a hack: We schedule complete_setup into the
     # event loop and return immediately (return True)
@@ -291,7 +291,7 @@ async def async_setup_entry(hass: HomeAssistantType,
 
 async def _setup_auto_reconnect_logic(hass: HomeAssistantType,
                                       cli: 'APIClient',
-                                      entry: ConfigEntry, host: str):
+                                      entry: ConfigEntry, host: str, on_login):
     """Set up the re-connect logic for the API client."""
     from aioesphomeapi import APIConnectionError
 
@@ -308,33 +308,40 @@ async def _setup_auto_reconnect_logic(hass: HomeAssistantType,
         data.available = False
         data.async_update_device_state(hass)
 
-        if tries != 0:
+        if is_disconnect:
-            # If not first re-try, wait and print message
-            wait_time = min(2**tries, 300)
-            _LOGGER.info("Trying to reconnect in %s seconds", wait_time)
-            await asyncio.sleep(wait_time)
-
-        if is_disconnect and tries == 0:
             # This can happen often depending on WiFi signal strength.
             # So therefore all these connection warnings are logged
             # as infos. The "unavailable" logic will still trigger so the
             # user knows if the device is not connected.
-            _LOGGER.info("Disconnected from API")
+            _LOGGER.info("Disconnected from ESPHome API for %s", host)
+
+        if tries != 0:
+            # If not first re-try, wait and print message
+            # Cap wait time at 1 minute. This is because while working on the
+            # device (e.g. soldering stuff), users don't want to have to wait
+            # a long time for their device to show up in HA again (this was
+            # mentioned a lot in early feedback)
+            #
+            # In the future another API will be set up so that the ESP can
+            # notify HA of connectivity directly, but for new we'll use a
+            # really short reconnect interval.
+            wait_time = int(round(min(1.8**tries, 60.0)))
+            _LOGGER.info("Trying to reconnect in %s seconds", wait_time)
+            await asyncio.sleep(wait_time)
 
         try:
-            await cli.connect()
+            await cli.connect(on_stop=try_connect, login=True)
-            await cli.login()
         except APIConnectionError as error:
-            _LOGGER.info("Can't connect to esphome API for '%s' (%s)",
+            _LOGGER.info("Can't connect to ESPHome API for %s: %s",
                          host, error)
             # Schedule re-connect in event loop in order not to delay HA
             # startup. First connect is scheduled in tracked tasks.
-            data.reconnect_task = \
+            data.reconnect_task = hass.loop.create_task(
-                hass.loop.create_task(try_connect(tries + 1, is_disconnect))
+                try_connect(tries + 1, is_disconnect=False))
         else:
             _LOGGER.info("Successfully connected to %s", host)
+            hass.async_create_task(on_login())
 
-    cli.on_disconnect = try_connect
     return try_connect
 
 
@@ -368,7 +375,7 @@ async def _cleanup_instance(hass: HomeAssistantType,
         disconnect_cb()
     for cleanup_callback in data.cleanup_callbacks:
         cleanup_callback()
-    await data.client.stop()
+    await data.client.disconnect()
 
 
 async def async_unload_entry(hass: HomeAssistantType,

homeassistant/components/esphome/config_flow.py

@@ -92,7 +92,6 @@ class EsphomeFlowHandler(config_entries.ConfigFlow):
         cli = APIClient(self.hass.loop, self._host, self._port, '')
 
         try:
-            await cli.start()
             await cli.connect()
             device_info = await cli.device_info()
         except APIConnectionError as err:
@@ -100,7 +99,7 @@ class EsphomeFlowHandler(config_entries.ConfigFlow):
                 return 'resolve_error', None
             return 'connection_error', None
         finally:
-            await cli.stop(force=True)
+            await cli.disconnect(force=True)
 
         return None, device_info
 
@@ -111,17 +110,9 @@ class EsphomeFlowHandler(config_entries.ConfigFlow):
         cli = APIClient(self.hass.loop, self._host, self._port, self._password)
 
         try:
-            await cli.start()
+            await cli.connect(login=True)
-            await cli.connect()
-        except APIConnectionError:
-            await cli.stop(force=True)
-            return 'connection_error'
-
-        try:
-            await cli.login()
         except APIConnectionError:
+            await cli.disconnect(force=True)
             return 'invalid_password'
-        finally:
-            await cli.stop(force=True)
 
         return None

requirements_all.txt

@@ -96,7 +96,7 @@ aioautomatic==0.6.5
 aiodns==1.1.1
 
 # homeassistant.components.esphome
-aioesphomeapi==1.3.0
+aioesphomeapi==1.4.0
 
 # homeassistant.components.freebox
 aiofreepybox==0.0.6

tests/components/esphome/test_config_flow.py

@@ -31,10 +31,8 @@ def mock_client():
             return mock_client
 
         mock_client.side_effect = mock_constructor
-        mock_client.start.return_value = mock_coro()
         mock_client.connect.return_value = mock_coro()
-        mock_client.stop.return_value = mock_coro()
+        mock_client.disconnect.return_value = mock_coro()
-        mock_client.login.return_value = mock_coro()
 
         yield mock_client
 
@@ -69,10 +67,9 @@ async def test_user_connection_works(hass, mock_client):
         'password': ''
     }
     assert result['title'] == 'test'
-    assert len(mock_client.start.mock_calls) == 1
     assert len(mock_client.connect.mock_calls) == 1
     assert len(mock_client.device_info.mock_calls) == 1
-    assert len(mock_client.stop.mock_calls) == 1
+    assert len(mock_client.disconnect.mock_calls) == 1
     assert mock_client.host == '127.0.0.1'
     assert mock_client.port == 80
     assert mock_client.password == ''
@@ -106,10 +103,9 @@ async def test_user_resolve_error(hass, mock_api_connection_error,
     assert result['errors'] == {
         'base': 'resolve_error'
     }
-    assert len(mock_client.start.mock_calls) == 1
     assert len(mock_client.connect.mock_calls) == 1
     assert len(mock_client.device_info.mock_calls) == 1
-    assert len(mock_client.stop.mock_calls) == 1
+    assert len(mock_client.disconnect.mock_calls) == 1
 
 
 async def test_user_connection_error(hass, mock_api_connection_error,
@@ -131,10 +127,9 @@ async def test_user_connection_error(hass, mock_api_connection_error,
     assert result['errors'] == {
         'base': 'connection_error'
     }
-    assert len(mock_client.start.mock_calls) == 1
     assert len(mock_client.connect.mock_calls) == 1
     assert len(mock_client.device_info.mock_calls) == 1
-    assert len(mock_client.stop.mock_calls) == 1
+    assert len(mock_client.disconnect.mock_calls) == 1
 
 
 async def test_user_with_password(hass, mock_client):
@@ -176,12 +171,12 @@ async def test_user_invalid_password(hass, mock_api_connection_error,
 
     mock_client.device_info.return_value = mock_coro(
         MockDeviceInfo(True, "test"))
-    mock_client.login.side_effect = mock_api_connection_error
 
     await flow.async_step_user(user_input={
         'host': '127.0.0.1',
         'port': 6053,
     })
+    mock_client.connect.side_effect = mock_api_connection_error
     result = await flow.async_step_authenticate(user_input={
         'password': 'invalid'
     })
@@ -191,30 +186,3 @@ async def test_user_invalid_password(hass, mock_api_connection_error,
     assert result['errors'] == {
         'base': 'invalid_password'
     }
-
-
-async def test_user_login_connection_error(hass, mock_api_connection_error,
-                                           mock_client):
-    """Test user step with connection error during login phase."""
-    flow = config_flow.EsphomeFlowHandler()
-    flow.hass = hass
-    await flow.async_step_user(user_input=None)
-
-    mock_client.device_info.return_value = mock_coro(
-        MockDeviceInfo(True, "test"))
-
-    await flow.async_step_user(user_input={
-        'host': '127.0.0.1',
-        'port': 6053,
-    })
-
-    mock_client.connect.side_effect = mock_api_connection_error
-    result = await flow.async_step_authenticate(user_input={
-        'password': 'invalid'
-    })
-
-    assert result['type'] == 'form'
-    assert result['step_id'] == 'authenticate'
-    assert result['errors'] == {
-        'base': 'connection_error'
-    }