From 3978d0a1d0515b3c6cca12f71569f6b2f828ddc4 Mon Sep 17 00:00:00 2001
From: cyr-ius <cyr-ius@ipocus.net>
Date: Sat, 28 Mar 2020 18:06:08 +0100
Subject: [PATCH] Add async_connect

---
 homeassistant/components/roomba/__init__.py | 79 ++++++++++++++-------
 homeassistant/components/roomba/const.py    |  2 +-
 2 files changed, 56 insertions(+), 25 deletions(-)

diff --git a/homeassistant/components/roomba/__init__.py b/homeassistant/components/roomba/__init__.py
index 2efc16df425..c95115d5707 100644
--- a/homeassistant/components/roomba/__init__.py
+++ b/homeassistant/components/roomba/__init__.py
@@ -1,13 +1,14 @@
 """The roomba component."""
+import asyncio
 import logging
 
 import async_timeout
-from roomba import Roomba, RoombaConnectionError
+from roomba import Roomba
 
 from homeassistant import config_entries
-from homeassistant.const import CONF_HOST, CONF_NAME, CONF_PASSWORD, CONF_USERNAME
+from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME
 
-from .const import CONF_CERT, CONF_CONTINUOUS, CONF_DELAY, DOMAIN
+from .const import COMPONENTS, CONF_CERT, CONF_CONTINUOUS, CONF_DELAY, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -44,33 +45,63 @@ async def async_setup_entry(hass, config_entry):
     if DOMAIN not in hass.data:
         hass.data[DOMAIN] = {}
 
-    if "roomba" not in hass.data[DOMAIN]:
-        roomba = Roomba(
-            address=config_entry.data[CONF_HOST],
-            blid=config_entry.data[CONF_USERNAME],
-            password=config_entry.data[CONF_PASSWORD],
-            cert_name=config_entry.options[CONF_CERT],
-            continuous=config_entry.options[CONF_CONTINUOUS],
-            delay=config_entry.options[CONF_DELAY],
-        )
-        hass.data[DOMAIN]["roomba"] = roomba
-        try:
-            with async_timeout.timeout(9):
-                await hass.async_add_job(roomba.connect)
-        except RoombaConnectionError:
-            _LOGGER.error("Error to connect to %s", config_entry.data[CONF_HOST])
-            return False
+    config_entry.add_update_listener(async_connect)
+    if not await async_connect(hass, config_entry):
+        return False
+
+    for component in COMPONENTS:
+        hass.async_create_task(
+            hass.config_entries.async_forward_entry_setup(config_entry, component)
+        )
 
-    hass.data[DOMAIN]["name"] = config_entry.data[CONF_NAME]
-    hass.async_create_task(
-        hass.config_entries.async_forward_entry_setup(config_entry, "vacuum")
-    )
     return True
 
 
 async def async_unload_entry(hass, config_entry):
     """Unload a config entry."""
-    await hass.config_entries.async_forward_entry_unload(config_entry, "vacuum")
+    for component in COMPONENTS:
+        hass.async_create_task(
+            hass.config_entries.async_forward_entry_unload(config_entry, component)
+        )
     roomba = hass.data[DOMAIN]["roomba"]
     await hass.async_add_job(roomba.disconnect)
     return True
+
+
+async def async_connect(hass, config_entry):
+    """Connect to vacuum."""
+    # Check if triggred listerner
+    if "roomba" in hass.data[DOMAIN]:
+        await hass.async_add_job(hass.data[DOMAIN]["roomba"].disconnect)
+        await asyncio.sleep(1)
+
+    roomba = Roomba(
+        address=config_entry.data[CONF_HOST],
+        blid=config_entry.data[CONF_USERNAME],
+        password=config_entry.data[CONF_PASSWORD],
+        cert_name=config_entry.options[CONF_CERT],
+        continuous=config_entry.options[CONF_CONTINUOUS],
+        delay=config_entry.options[CONF_DELAY],
+    )
+
+    hass.data[DOMAIN]["roomba"] = roomba
+
+    try:
+        name = None
+        with async_timeout.timeout(9):
+            await hass.async_add_job(roomba.connect)
+            while not roomba.roomba_connected or name is None:
+                # Waiting for connection and check datas ready
+                name = (
+                    roomba.master_state.get("state", {})
+                    .get("reported", {})
+                    .get("name", None)
+                )
+                await asyncio.sleep(0.5)
+    except asyncio.TimeoutError:
+        # api looping if user or password incorrect and roomba exist
+        _LOGGER.error("Timeout exceeded")
+        await hass.async_add_job(roomba.disconnect)
+        return False
+
+    return True
diff --git a/homeassistant/components/roomba/const.py b/homeassistant/components/roomba/const.py
index f5c1a1a3d3e..803a1ad87ff 100644
--- a/homeassistant/components/roomba/const.py
+++ b/homeassistant/components/roomba/const.py
@@ -1,9 +1,9 @@
 """The roomba constants."""
 DOMAIN = "roomba"
+COMPONENTS = ["sensor", "binary_sensor", "vacuum"]
 CONF_CERT = "certificate"
 CONF_CONTINUOUS = "continuous"
 CONF_DELAY = "delay"
 DEFAULT_CERT = "/etc/ssl/certs/ca-certificates.crt"
 DEFAULT_CONTINUOUS = True
 DEFAULT_DELAY = 1
-DEFAULT_NAME = "Roomba"