"""Support for Neato botvac connected vacuum cleaners.""" import asyncio import logging from datetime import timedelta from pybotvac.exceptions import NeatoException, NeatoLoginException, NeatoRobotException import voluptuous as vol from homeassistant.config_entries import SOURCE_IMPORT from homeassistant.const import CONF_PASSWORD, CONF_USERNAME from homeassistant.helpers import config_validation as cv from homeassistant.util import Throttle from .config_flow import NeatoConfigFlow from .const import ( CONF_VENDOR, NEATO_CONFIG, NEATO_DOMAIN, NEATO_LOGIN, NEATO_ROBOTS, NEATO_PERSISTENT_MAPS, NEATO_MAP_DATA, SCAN_INTERVAL_MINUTES, VALID_VENDORS, ) _LOGGER = logging.getLogger(__name__) CONFIG_SCHEMA = vol.Schema( { NEATO_DOMAIN: vol.Schema( { vol.Required(CONF_USERNAME): cv.string, vol.Required(CONF_PASSWORD): cv.string, vol.Optional(CONF_VENDOR, default="neato"): vol.In(VALID_VENDORS), } ) }, extra=vol.ALLOW_EXTRA, ) async def async_setup(hass, config): """Set up the Neato component.""" if NEATO_DOMAIN not in config: # There is an entry and nothing in configuration.yaml return True entries = hass.config_entries.async_entries(NEATO_DOMAIN) hass.data[NEATO_CONFIG] = config[NEATO_DOMAIN] if entries: # There is an entry and something in the configuration.yaml entry = entries[0] conf = config[NEATO_DOMAIN] if ( entry.data[CONF_USERNAME] == conf[CONF_USERNAME] and entry.data[CONF_PASSWORD] == conf[CONF_PASSWORD] and entry.data[CONF_VENDOR] == conf[CONF_VENDOR] ): # The entry is not outdated return True # The entry is outdated error = await hass.async_add_executor_job( NeatoConfigFlow.try_login, conf[CONF_USERNAME], conf[CONF_PASSWORD], conf[CONF_VENDOR], ) if error is not None: _LOGGER.error(error) return False # Update the entry hass.config_entries.async_update_entry(entry, data=config[NEATO_DOMAIN]) else: # Create the new entry hass.async_create_task( hass.config_entries.flow.async_init( NEATO_DOMAIN, context={"source": SOURCE_IMPORT}, data=config[NEATO_DOMAIN], ) ) return True async def async_setup_entry(hass, entry): """Set up config entry.""" from pybotvac import Account, Neato, Vorwerk if entry.data[CONF_VENDOR] == "neato": hass.data[NEATO_LOGIN] = NeatoHub(hass, entry.data, Account, Neato) elif entry.data[CONF_VENDOR] == "vorwerk": hass.data[NEATO_LOGIN] = NeatoHub(hass, entry.data, Account, Vorwerk) hub = hass.data[NEATO_LOGIN] await hass.async_add_executor_job(hub.login) if not hub.logged_in: _LOGGER.debug("Failed to login to Neato API") return False try: await hass.async_add_executor_job(hub.update_robots) except NeatoRobotException: _LOGGER.debug("Failed to connect to Neato API") return False for component in ("camera", "vacuum", "switch", "sensor"): hass.async_create_task( hass.config_entries.async_forward_entry_setup(entry, component) ) return True async def async_unload_entry(hass, entry): """Unload config entry.""" hass.data.pop(NEATO_LOGIN) await asyncio.gather( hass.config_entries.async_forward_entry_unload(entry, "camera"), hass.config_entries.async_forward_entry_unload(entry, "vacuum"), hass.config_entries.async_forward_entry_unload(entry, "switch"), hass.config_entries.async_forward_entry_unload(entry, "sensor"), ) return True class NeatoHub: """A My Neato hub wrapper class.""" def __init__(self, hass, domain_config, neato, vendor): """Initialize the Neato hub.""" self.config = domain_config self._neato = neato self._hass = hass self._vendor = vendor self.my_neato = None self.logged_in = False def login(self): """Login to My Neato.""" _LOGGER.debug("Trying to connect to Neato API") try: self.my_neato = self._neato( self.config[CONF_USERNAME], self.config[CONF_PASSWORD], self._vendor ) except NeatoException as ex: if isinstance(ex, NeatoLoginException): _LOGGER.error("Invalid credentials") else: _LOGGER.error("Unable to connect to Neato API") self.logged_in = False return self.logged_in = True _LOGGER.debug("Successfully connected to Neato API") @Throttle(timedelta(minutes=SCAN_INTERVAL_MINUTES)) def update_robots(self): """Update the robot states.""" _LOGGER.debug("Running HUB.update_robots %s", self._hass.data.get(NEATO_ROBOTS)) self._hass.data[NEATO_ROBOTS] = self.my_neato.robots self._hass.data[NEATO_PERSISTENT_MAPS] = self.my_neato.persistent_maps self._hass.data[NEATO_MAP_DATA] = self.my_neato.maps def download_map(self, url): """Download a new map image.""" map_image_data = self.my_neato.get_map_image(url) return map_image_data