Add connect and disconnect functions

homeassistant/components/roomba/__init__.py

@@ -3,9 +3,9 @@ import asyncio
 import logging
 
 import async_timeout
-from roomba import Roomba
+from roomba import Roomba, RoombaConnectionError
 
-from homeassistant import config_entries
+from homeassistant import config_entries, exceptions
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME
 
 from .const import COMPONENTS, CONF_CERT, CONF_CONTINUOUS, CONF_DELAY, DOMAIN
@@ -45,35 +45,7 @@ async def async_setup_entry(hass, config_entry):
     if DOMAIN not in hass.data:
         hass.data[DOMAIN] = {}
 
-    config_entry.add_update_listener(async_connect)
+    if "roomba" not in hass.data[DOMAIN]:
-    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)
-        )
-
-    return True
-
-
-async def async_unload_entry(hass, config_entry):
-    """Unload a config entry."""
-    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 triggered listener
-    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],
@@ -84,9 +56,24 @@ async def async_connect(hass, config_entry):
             delay=config_entry.options[CONF_DELAY],
         )
 
+        if not await async_connect_or_timeout(hass, roomba):
+            return False
+
+    for component in COMPONENTS:
+        hass.async_create_task(
+            hass.config_entries.async_forward_entry_setup(config_entry, component)
+        )
+
+    config_entry.add_update_listener(async_update_options)
+
+    return True
+
+
+async def async_connect_or_timeout(hass, roomba):
+    """Connect to vacuum."""
     try:
-        name = None
+        hass.data[DOMAIN]["name"] = name = None
-        with async_timeout.timeout(9):
+        with async_timeout.timeout(10):
             await hass.async_add_job(roomba.connect)
             while not roomba.roomba_connected or name is None:
                 # Waiting for connection and check datas ready
@@ -97,10 +84,46 @@ async def async_connect(hass, config_entry):
                 )
                 await asyncio.sleep(0.5)
                 hass.data[DOMAIN]["roomba"] = roomba
+                hass.data[DOMAIN]["name"] = name
+    except RoombaConnectionError:
+        _LOGGER.error("Error to connect to vacuum")
+        raise CannotConnect
     except asyncio.TimeoutError:
         # api looping if user or password incorrect and roomba exist
-        _LOGGER.error("Timeout exceeded")
+        await async_disconnect_or_timeout(hass, roomba)
-        await hass.async_add_job(roomba.disconnect)
+        _LOGGER.exception("Timeout expired, user or password incorrect")
-        return False
+        raise InvalidAuth
 
     return True
+
+
+async def async_disconnect_or_timeout(hass, roomba):
+    """Disconnect to vacuum."""
+    await hass.async_add_job(roomba.disconnect)
+    await asyncio.sleep(1)
+    return True
+
+
+async def async_update_options(hass, config_entry):
+    """Update options."""
+    roomba = hass.data[DOMAIN]["roomba"]
+    await async_disconnect_or_timeout(hass, roomba)
+    await async_connect_or_timeout(hass, roomba)
+
+
+async def async_unload_entry(hass, config_entry):
+    """Unload a config entry."""
+    for component in COMPONENTS:
+        hass.async_create_task(
+            hass.config_entries.async_forward_entry_unload(config_entry, component)
+        )
+    roomba = hass.data[DOMAIN]["roomba"]
+    return await async_disconnect_or_timeout(hass, roomba)
+
+
+class CannotConnect(exceptions.HomeAssistantError):
+    """Error to indicate we cannot connect."""
+
+
+class InvalidAuth(exceptions.HomeAssistantError):
+    """Error to indicate there is invalid auth."""

homeassistant/components/roomba/config_flow.py

@@ -1,15 +1,14 @@
 """Config flow to configure demo component."""
-import asyncio
 import logging
 
-import async_timeout
+from roomba import Roomba
-from roomba import Roomba, RoombaConnectionError
 import voluptuous as vol
 
-from homeassistant import config_entries, core, exceptions
+from homeassistant import config_entries, core
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import callback
 
+from . import CannotConnect, InvalidAuth, async_connect_or_timeout
 from .const import (
     CONF_CERT,
     CONF_CONTINUOUS,
@@ -39,7 +38,7 @@ async def validate_input(hass: core.HomeAssistant, data):
 
     Data has the keys from DATA_SCHEMA with values provided by the user.
     """
-    api = Roomba(
+    roomba = Roomba(
         address=data[CONF_HOST],
         blid=data[CONF_USERNAME],
         password=data[CONF_PASSWORD],
@@ -47,26 +46,8 @@ async def validate_input(hass: core.HomeAssistant, data):
         continuous=data[CONF_CONTINUOUS],
         delay=data[CONF_DELAY],
     )
-    try:
-        name = None
-        with async_timeout.timeout(10):
-            await hass.async_add_job(api.connect)
-            while not api.roomba_connected or name is None:
-                name = (
-                    api.master_state.get("state", {})
-                    .get("reported", {})
-                    .get("name", None)
-                )
-                await asyncio.sleep(0.5)
-            hass.data[DOMAIN]["roomba"] = api
-    except RoombaConnectionError:
-        raise CannotConnect
-    except asyncio.TimeoutError:
-        # Api looping if user or password incorrect and roomba exist
-        await hass.async_add_job(api.disconnect)
-        raise InvalidAuth
 
-    return {"title": name}
+    await async_connect_or_timeout(hass, roomba)
 
 
 class RoombaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
@@ -93,17 +74,18 @@ class RoombaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
         if user_input is not None:
             try:
-                info = await validate_input(self.hass, user_input)
+                await validate_input(self.hass, user_input)
             except CannotConnect:
                 errors = {"base": "cannot_connect"}
             except InvalidAuth:
                 errors = {"base": "invalid_auth"}
             except Exception:  # pylint: disable=broad-except
-                _LOGGER.exception("Unexpected exception")
+                errors = {"base": "unknown"}
-                errors["base"] = "unknown"
 
             if "base" not in errors:
-                return self.async_create_entry(title=info["title"], data=user_input)
+                return self.async_create_entry(
+                    title=self.hass.data[DOMAIN]["name"], data=user_input
+                )
 
         # If there was no user input, do not show the errors.
         if user_input is None:
@@ -149,11 +131,3 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
                 }
             ),
         )
-
-
-class CannotConnect(exceptions.HomeAssistantError):
-    """Error to indicate we cannot connect."""
-
-
-class InvalidAuth(exceptions.HomeAssistantError):
-    """Error to indicate there is invalid auth."""