diff --git a/homeassistant/components/roomba/__init__.py b/homeassistant/components/roomba/__init__.py index a933b99bcec..10eb24a4f1f 100644 --- a/homeassistant/components/roomba/__init__.py +++ b/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,48 +45,35 @@ 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 not await async_connect(hass, config_entry): - return False + 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], + ) + + 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_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): +async def async_connect_or_timeout(hass, roomba): """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], - 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], - ) - try: - name = None - with async_timeout.timeout(9): + hass.data[DOMAIN]["name"] = name = None + 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 @@ -96,11 +83,47 @@ async def async_connect(hass, config_entry): .get("name", None) ) await asyncio.sleep(0.5) - hass.data[DOMAIN]["roomba"] = roomba + 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 hass.async_add_job(roomba.disconnect) - return False + await async_disconnect_or_timeout(hass, roomba) + _LOGGER.exception("Timeout expired, user or password incorrect") + 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.""" diff --git a/homeassistant/components/roomba/config_flow.py b/homeassistant/components/roomba/config_flow.py index 9fde62e3c3b..4324e21dba4 100644 --- a/homeassistant/components/roomba/config_flow.py +++ b/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, RoombaConnectionError +from roomba import Roomba 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."""