core/homeassistant/components/neato/__init__.py

"""Support for Neato botvac connected vacuum cleaners."""
import asyncio
from datetime import timedelta
import logging

from pybotvac import Account, Neato, Vorwerk
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.exceptions import ConfigEntryNotReady
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,
    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."""
    hub = NeatoHub(hass, entry.data, Account)

    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 as ex:
        _LOGGER.debug("Failed to connect to Neato API")
        raise ConfigEntryNotReady from ex

    hass.data[NEATO_LOGIN] = hub

    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")
                raise ConfigEntryNotReady from ex
            self.logged_in = False
            return

        self.logged_in = True
        _LOGGER.debug("Successfully connected to Neato API")

    @Throttle(timedelta(minutes=1))
    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