"""Support for Neato botvac connected vacuum cleaners.""" import asyncio import logging from datetime import timedelta import voluptuous as vol from pybotvac import Account, Neato, Vorwerk from pybotvac.exceptions import NeatoException, NeatoLoginException, NeatoRobotException 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_MAP_DATA, NEATO_PERSISTENT_MAPS, NEATO_ROBOTS, 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.""" hass.data[NEATO_LOGIN] = NeatoHub(hass, entry.data, Account) 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): """Initialize the Neato hub.""" self.config = domain_config self._neato = neato self._hass = hass if self.config[CONF_VENDOR] == "vorwerk": self._vendor = Vorwerk() else: # Neato self._vendor = Neato() 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