"""Shark IQ Integration.""" import asyncio import async_timeout from sharkiqpy import ( AylaApi, SharkIqAuthError, SharkIqAuthExpiringError, SharkIqNotAuthedError, get_ayla_api, ) import voluptuous as vol from homeassistant import config_entries, exceptions from homeassistant.const import CONF_PASSWORD, CONF_USERNAME from .const import API_TIMEOUT, COMPONENTS, DOMAIN, LOGGER from .update_coordinator import SharkIqUpdateCoordinator CONFIG_SCHEMA = vol.Schema({DOMAIN: vol.Schema({})}, extra=vol.ALLOW_EXTRA) class CannotConnect(exceptions.HomeAssistantError): """Error to indicate we cannot connect.""" async def async_setup(hass, config): """Set up the sharkiq environment.""" hass.data.setdefault(DOMAIN, {}) if DOMAIN not in config: return True for index, conf in enumerate(config[DOMAIN]): LOGGER.debug( "Importing Shark IQ #%d (Username: %s)", index, conf[CONF_USERNAME] ) hass.async_create_task( hass.config_entries.flow.async_init( DOMAIN, context={"source": config_entries.SOURCE_IMPORT}, data=conf, ) ) async def async_connect_or_timeout(ayla_api: AylaApi) -> AylaApi: """Connect to vacuum.""" try: with async_timeout.timeout(API_TIMEOUT): LOGGER.debug("Initialize connection to Ayla networks API") await ayla_api.async_sign_in() except SharkIqAuthError as exc: LOGGER.error("Error to connect to Shark IQ api") raise CannotConnect from exc except asyncio.TimeoutError as exc: LOGGER.error("Timeout expired") raise CannotConnect from exc return ayla_api async def async_setup_entry(hass, config_entry): """Initialize the sharkiq platform via config entry.""" ayla_api = get_ayla_api( username=config_entry.data[CONF_USERNAME], password=config_entry.data[CONF_PASSWORD], websession=hass.helpers.aiohttp_client.async_get_clientsession(), ) try: if not await async_connect_or_timeout(ayla_api): return False except CannotConnect as exc: raise exceptions.ConfigEntryNotReady from exc shark_vacs = await ayla_api.async_get_devices(False) device_names = ", ".join([d.name for d in shark_vacs]) LOGGER.info("Found %d Shark IQ device(s): %s", len(device_names), device_names) coordinator = SharkIqUpdateCoordinator(hass, config_entry, ayla_api, shark_vacs) await coordinator.async_refresh() if not coordinator.last_update_success: raise exceptions.ConfigEntryNotReady hass.data[DOMAIN][config_entry.entry_id] = coordinator for component in COMPONENTS: hass.async_create_task( hass.config_entries.async_forward_entry_setup(config_entry, component) ) if not config_entry.update_listeners: config_entry.add_update_listener(async_update_options) return True async def async_disconnect_or_timeout(coordinator: SharkIqUpdateCoordinator): """Disconnect to vacuum.""" LOGGER.debug("Disconnecting from Ayla Api") with async_timeout.timeout(5): try: await coordinator.ayla_api.async_sign_out() except (SharkIqAuthError, SharkIqAuthExpiringError, SharkIqNotAuthedError): pass return True async def async_update_options(hass, config_entry): """Update options.""" await hass.config_entries.async_reload(config_entry.entry_id) async def async_unload_entry(hass, config_entry): """Unload a config entry.""" unload_ok = all( await asyncio.gather( *[ hass.config_entries.async_forward_entry_unload(config_entry, component) for component in COMPONENTS ] ) ) if unload_ok: domain_data = hass.data[DOMAIN][config_entry.entry_id] try: await async_disconnect_or_timeout(coordinator=domain_data) except SharkIqAuthError: pass hass.data[DOMAIN].pop(config_entry.entry_id) return unload_ok