From a7950937058903808d87256537d40d3b0b7c93f2 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Tue, 16 Oct 2018 10:35:35 +0200 Subject: [PATCH] UniFi POE control (#17011) * First commit * Feature complete? * Add dependency * Move setting poe mode logic to library * Use guard clauses * Bump requirement to 2 * Simplify saving switches with poe off * Store and use poe mode * Fix indentation * Fix flake8 * Configuration future proofing * Bump dependency to v3 * Add first test * Proper use of defaults with config flow (thanks helto) * Appease hound * Make sure there can't be duplicate entries of combination host+site * More tests * More tests * 98% coverage of controller * Fix hound comments * Config flow step init not necessary * Use async_current_entries to check if host and site for controller is used * Remove storing/restoring poe off devices to slim PR * First batch of switch tests * More switch tests. * Small improvements and clean up * Make tests pass Don't name device in device registry * Dont process clients that belong to non-UniFi POE switches * Allow selection of site from a list in config flow * Fix double blank lines in method * Update codeowners --- CODEOWNERS | 2 + homeassistant/components/switch/unifi.py | 230 ++++++++++++ .../components/unifi/.translations/en.json | 26 ++ homeassistant/components/unifi/__init__.py | 186 ++++++++++ homeassistant/components/unifi/const.py | 12 + homeassistant/components/unifi/controller.py | 131 +++++++ homeassistant/components/unifi/errors.py | 26 ++ homeassistant/components/unifi/strings.json | 26 ++ homeassistant/config_entries.py | 1 + requirements_all.txt | 3 + requirements_test_all.txt | 3 + script/gen_requirements_all.py | 1 + tests/components/switch/test_unifi.py | 345 ++++++++++++++++++ tests/components/unifi/__init__.py | 1 + tests/components/unifi/test_controller.py | 266 ++++++++++++++ tests/components/unifi/test_init.py | 330 +++++++++++++++++ 16 files changed, 1589 insertions(+) create mode 100644 homeassistant/components/switch/unifi.py create mode 100644 homeassistant/components/unifi/.translations/en.json create mode 100644 homeassistant/components/unifi/__init__.py create mode 100644 homeassistant/components/unifi/const.py create mode 100644 homeassistant/components/unifi/controller.py create mode 100644 homeassistant/components/unifi/errors.py create mode 100644 homeassistant/components/unifi/strings.json create mode 100644 tests/components/switch/test_unifi.py create mode 100644 tests/components/unifi/__init__.py create mode 100644 tests/components/unifi/test_controller.py create mode 100644 tests/components/unifi/test_init.py diff --git a/CODEOWNERS b/CODEOWNERS index 9343407f06f..40a33c66b78 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -224,6 +224,8 @@ homeassistant/components/tradfri/* @ggravlingen homeassistant/components/*/tradfri.py @ggravlingen # U +homeassistant/components/unifi.py @kane610 +homeassistant/components/switch/unifi.py @kane610 homeassistant/components/upcloud.py @scop homeassistant/components/*/upcloud.py @scop diff --git a/homeassistant/components/switch/unifi.py b/homeassistant/components/switch/unifi.py new file mode 100644 index 00000000000..dc02068c4a8 --- /dev/null +++ b/homeassistant/components/switch/unifi.py @@ -0,0 +1,230 @@ +""" +Support for devices connected to UniFi POE. + +For more details about this platform, please refer to the documentation at +https://home-assistant.io/components/switch.unifi/ +""" + +import asyncio +import logging + +from datetime import timedelta + +import async_timeout + +from homeassistant.components import unifi +from homeassistant.components.switch import SwitchDevice +from homeassistant.components.unifi.const import ( + CONF_CONTROLLER, CONF_SITE_ID, CONTROLLER_ID, DOMAIN) +from homeassistant.const import CONF_HOST +from homeassistant.core import callback +from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC + +DEPENDENCIES = [DOMAIN] +SCAN_INTERVAL = timedelta(seconds=15) + +LOGGER = logging.getLogger(__name__) + + +async def async_setup_platform( + hass, config, async_add_entities, discovery_info=None): + """Component doesn't support configuration through configuration.yaml.""" + pass + + +async def async_setup_entry(hass, config_entry, async_add_entities): + """Set up switches for UniFi component. + + Switches are controlling network switch ports with Poe. + """ + controller_id = CONTROLLER_ID.format( + host=config_entry.data[CONF_CONTROLLER][CONF_HOST], + site=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID] + ) + controller = hass.data[unifi.DOMAIN][controller_id] + switches = {} + + progress = None + update_progress = set() + + async def request_update(object_id): + """Request an update.""" + nonlocal progress + update_progress.add(object_id) + + if progress is not None: + return await progress + + progress = asyncio.ensure_future(update_controller()) + result = await progress + progress = None + update_progress.clear() + return result + + async def update_controller(): + """Update the values of the controller.""" + tasks = [async_update_items( + controller, async_add_entities, request_update, + switches, update_progress + )] + await asyncio.wait(tasks) + + await update_controller() + + +async def async_update_items(controller, async_add_entities, + request_controller_update, switches, + progress_waiting): + """Update POE port state from the controller.""" + import aiounifi + + @callback + def update_switch_state(): + """Tell switches to reload state.""" + for client_id, client in switches.items(): + if client_id not in progress_waiting: + client.async_schedule_update_ha_state() + + try: + with async_timeout.timeout(4): + await controller.api.clients.update() + await controller.api.devices.update() + + except aiounifi.LoginRequired: + try: + with async_timeout.timeout(5): + await controller.api.login() + except (asyncio.TimeoutError, aiounifi.AiounifiException): + if controller.available: + controller.available = False + update_switch_state() + return + + except (asyncio.TimeoutError, aiounifi.AiounifiException): + if controller.available: + LOGGER.error('Unable to reach controller %s', controller.host) + controller.available = False + update_switch_state() + return + + if not controller.available: + LOGGER.info('Reconnected to controller %s', controller.host) + controller.available = True + + new_switches = [] + devices = controller.api.devices + for client_id in controller.api.clients: + + if client_id in progress_waiting: + continue + + if client_id in switches: + LOGGER.debug("Updating UniFi switch %s (%s)", + switches[client_id].entity_id, + switches[client_id].client.mac) + switches[client_id].async_schedule_update_ha_state() + continue + + client = controller.api.clients[client_id] + # Network device with active POE + if not client.is_wired or client.sw_mac not in devices or \ + not devices[client.sw_mac].ports[client.sw_port].port_poe or \ + not devices[client.sw_mac].ports[client.sw_port].poe_enable: + continue + + # Multiple POE-devices on same port means non UniFi POE driven switch + multi_clients_on_port = False + for client2 in controller.api.clients.values(): + if client.mac != client2.mac and \ + client.sw_mac == client2.sw_mac and \ + client.sw_port == client2.sw_port: + multi_clients_on_port = True + break + + if multi_clients_on_port: + continue + + switches[client_id] = UniFiSwitch( + client, controller, request_controller_update) + new_switches.append(switches[client_id]) + LOGGER.debug("New UniFi switch %s (%s)", client.hostname, client.mac) + + if new_switches: + async_add_entities(new_switches) + + +class UniFiSwitch(SwitchDevice): + """Representation of a client that uses POE.""" + + def __init__(self, client, controller, request_controller_update): + """Set up switch.""" + self.client = client + self.controller = controller + self.poe_mode = None + if self.port.poe_mode != 'off': + self.poe_mode = self.port.poe_mode + self.async_request_controller_update = request_controller_update + + async def async_update(self): + """Synchronize state with controller.""" + await self.async_request_controller_update(self.client.mac) + + @property + def name(self): + """Return the name of the switch.""" + return self.client.hostname + + @property + def unique_id(self): + """Return a unique identifier for this switch.""" + return 'poe-{}'.format(self.client.mac) + + @property + def is_on(self): + """Return true if POE is active.""" + return self.port.poe_mode != 'off' + + @property + def available(self): + """Return if switch is available.""" + return self.controller.available or \ + self.client.sw_mac in self.controller.api.devices + + async def async_turn_on(self, **kwargs): + """Enable POE for client.""" + await self.device.async_set_port_poe_mode( + self.client.sw_port, self.poe_mode) + + async def async_turn_off(self, **kwargs): + """Disable POE for client.""" + await self.device.async_set_port_poe_mode(self.client.sw_port, 'off') + + @property + def device_state_attributes(self): + """Return the device state attributes.""" + attributes = { + 'power': self.port.poe_power, + 'received': self.client.wired_rx_bytes / 1000000, + 'sent': self.client.wired_tx_bytes / 1000000, + 'switch': self.client.sw_mac, + 'port': self.client.sw_port, + 'poe_mode': self.poe_mode + } + return attributes + + @property + def device_info(self): + """Return a device description for device registry.""" + return { + 'connections': {(CONNECTION_NETWORK_MAC, self.client.mac)} + } + + @property + def device(self): + """Shortcut to the switch that client is connected to.""" + return self.controller.api.devices[self.client.sw_mac] + + @property + def port(self): + """Shortcut to the switch port that client is connected to.""" + return self.device.ports[self.client.sw_port] diff --git a/homeassistant/components/unifi/.translations/en.json b/homeassistant/components/unifi/.translations/en.json new file mode 100644 index 00000000000..938ac058d22 --- /dev/null +++ b/homeassistant/components/unifi/.translations/en.json @@ -0,0 +1,26 @@ +{ + "config": { + "title": "UniFi Controller", + "step": { + "user": { + "title": "Set up UniFi Controller", + "data": { + "host": "Host", + "username": "User name", + "password": "Password", + "port": "Port", + "site": "Site ID", + "verify_ssl": "Controller using proper certificate" + } + } + }, + "error": { + "faulty_credentials": "Bad user credentials", + "service_unavailable": "No service available" + }, + "abort": { + "already_configured": "Controller site is already configured", + "user_privilege": "User needs to be administrator" + } + } +} \ No newline at end of file diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py new file mode 100644 index 00000000000..26b60aecf42 --- /dev/null +++ b/homeassistant/components/unifi/__init__.py @@ -0,0 +1,186 @@ +""" +Support for devices connected to UniFi POE. + +For more details about this platform, please refer to the documentation at +https://home-assistant.io/components/unifi/ +""" + +import voluptuous as vol + +from homeassistant import config_entries +from homeassistant.const import ( + CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME, CONF_VERIFY_SSL) +from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC + +from .const import (CONF_CONTROLLER, CONF_POE_CONTROL, CONF_SITE_ID, + CONTROLLER_ID, DOMAIN, LOGGER) +from .controller import UniFiController, get_controller +from .errors import ( + AlreadyConfigured, AuthenticationRequired, CannotConnect, UserLevel) + +DEFAULT_PORT = 8443 +DEFAULT_SITE_ID = 'default' +DEFAULT_VERIFY_SSL = False + +REQUIREMENTS = ['aiounifi==3'] + + +async def async_setup(hass, config): + """Component doesn't support configuration through configuration.yaml.""" + return True + + +async def async_setup_entry(hass, config_entry): + """Set up the UniFi component.""" + controller = UniFiController(hass, config_entry) + + if DOMAIN not in hass.data: + hass.data[DOMAIN] = {} + controller_id = CONTROLLER_ID.format( + host=config_entry.data[CONF_CONTROLLER][CONF_HOST], + site=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID] + ) + + if not await controller.async_setup(): + return False + + hass.data[DOMAIN][controller_id] = controller + + if controller.mac is None: + return True + + device_registry = await \ + hass.helpers.device_registry.async_get_registry() + device_registry.async_get_or_create( + config_entry_id=config_entry.entry_id, + connections={(CONNECTION_NETWORK_MAC, controller.mac)}, + manufacturer='Ubiquiti', + model="UniFi Controller", + name="UniFi Controller", + # sw_version=config.raw['swversion'], + ) + + return True + + +async def async_unload_entry(hass, config_entry): + """Unload a config entry.""" + controller_id = CONTROLLER_ID.format( + host=config_entry.data[CONF_CONTROLLER][CONF_HOST], + site=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID] + ) + controller = hass.data[DOMAIN].pop(controller_id) + return await controller.async_reset() + + +@config_entries.HANDLERS.register(DOMAIN) +class UnifiFlowHandler(config_entries.ConfigFlow): + """Handle a UniFi config flow.""" + + VERSION = 1 + CONNECTION_CLASS = config_entries.CONN_CLASS_LOCAL_POLL + + def __init__(self): + """Initialize the UniFi flow.""" + self.config = None + self.desc = None + self.sites = None + + async def async_step_user(self, user_input=None): + """Handle a flow initialized by the user.""" + errors = {} + + if user_input is not None: + + try: + self.config = { + CONF_HOST: user_input[CONF_HOST], + CONF_USERNAME: user_input[CONF_USERNAME], + CONF_PASSWORD: user_input[CONF_PASSWORD], + CONF_PORT: user_input.get(CONF_PORT), + CONF_VERIFY_SSL: user_input.get(CONF_VERIFY_SSL), + CONF_SITE_ID: DEFAULT_SITE_ID, + } + controller = await get_controller(self.hass, **self.config) + + self.sites = await controller.sites() + + return await self.async_step_site() + + except AuthenticationRequired: + errors['base'] = 'faulty_credentials' + + except CannotConnect: + errors['base'] = 'service_unavailable' + + except Exception: # pylint: disable=broad-except + LOGGER.error( + 'Unknown error connecting with UniFi Controller at %s', + user_input[CONF_HOST]) + return self.async_abort(reason='unknown') + + return self.async_show_form( + step_id='user', + data_schema=vol.Schema({ + vol.Required(CONF_HOST): str, + vol.Required(CONF_USERNAME): str, + vol.Required(CONF_PASSWORD): str, + vol.Optional(CONF_PORT, default=DEFAULT_PORT): int, + vol.Optional( + CONF_VERIFY_SSL, default=DEFAULT_VERIFY_SSL): bool, + }), + errors=errors, + ) + + async def async_step_site(self, user_input=None): + """Select site to control.""" + errors = {} + + if user_input is not None: + + try: + desc = user_input.get(CONF_SITE_ID, self.desc) + for site in self.sites.values(): + if desc == site['desc']: + if site['role'] != 'admin': + raise UserLevel + self.config[CONF_SITE_ID] = site['name'] + break + + for entry in self._async_current_entries(): + controller = entry.data[CONF_CONTROLLER] + if controller[CONF_HOST] == self.config[CONF_HOST] and \ + controller[CONF_SITE_ID] == self.config[CONF_SITE_ID]: + raise AlreadyConfigured + + data = { + CONF_CONTROLLER: self.config, + CONF_POE_CONTROL: True + } + + return self.async_create_entry( + title=desc, + data=data + ) + + except AlreadyConfigured: + return self.async_abort(reason='already_configured') + + except UserLevel: + return self.async_abort(reason='user_privilege') + + if len(self.sites) == 1: + self.desc = next(iter(self.sites.values()))['desc'] + return await self.async_step_site(user_input={}) + + sites = [] + for site in self.sites.values(): + sites.append(site['desc']) + + return self.async_show_form( + step_id='site', + data_schema=vol.Schema({ + vol.Required(CONF_SITE_ID): vol.In(sites) + }), + errors=errors, + ) diff --git a/homeassistant/components/unifi/const.py b/homeassistant/components/unifi/const.py new file mode 100644 index 00000000000..7250feec799 --- /dev/null +++ b/homeassistant/components/unifi/const.py @@ -0,0 +1,12 @@ +"""Constants for the UniFi component.""" + +import logging + +LOGGER = logging.getLogger('homeassistant.components.unifi') +DOMAIN = 'unifi' + +CONTROLLER_ID = '{host}-{site}' + +CONF_CONTROLLER = 'controller' +CONF_POE_CONTROL = 'poe_control' +CONF_SITE_ID = 'site' diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py new file mode 100644 index 00000000000..9e21956536f --- /dev/null +++ b/homeassistant/components/unifi/controller.py @@ -0,0 +1,131 @@ +"""UniFi Controller abstraction.""" + +import asyncio +import async_timeout + +from aiohttp import CookieJar + +from homeassistant import config_entries +from homeassistant.const import CONF_HOST +from homeassistant.helpers import aiohttp_client + +from .const import CONF_CONTROLLER, CONF_POE_CONTROL, LOGGER +from .errors import AuthenticationRequired, CannotConnect + + +class UniFiController: + """Manages a single UniFi Controller.""" + + def __init__(self, hass, config_entry): + """Initialize the system.""" + self.hass = hass + self.config_entry = config_entry + self.available = True + self.api = None + self.progress = None + self._cancel_retry_setup = None + + @property + def host(self): + """Return the host of this controller.""" + return self.config_entry.data[CONF_CONTROLLER][CONF_HOST] + + @property + def mac(self): + """Return the mac address of this controller.""" + for client in self.api.clients.values(): + if self.host == client.ip: + return client.mac + return None + + async def async_setup(self, tries=0): + """Set up a UniFi controller.""" + hass = self.hass + + try: + self.api = await get_controller( + self.hass, **self.config_entry.data[CONF_CONTROLLER]) + await self.api.initialize() + + except CannotConnect: + retry_delay = 2 ** (tries + 1) + LOGGER.error("Error connecting to the UniFi controller. Retrying " + "in %d seconds", retry_delay) + + async def retry_setup(_now): + """Retry setup.""" + if await self.async_setup(tries + 1): + # This feels hacky, we should find a better way to do this + self.config_entry.state = config_entries.ENTRY_STATE_LOADED + + self._cancel_retry_setup = hass.helpers.event.async_call_later( + retry_delay, retry_setup) + + return False + + except Exception: # pylint: disable=broad-except + LOGGER.error( + 'Unknown error connecting with UniFi controller.') + return False + + if self.config_entry.data[CONF_POE_CONTROL]: + hass.async_create_task( + hass.config_entries.async_forward_entry_setup( + self.config_entry, 'switch')) + + return True + + async def async_reset(self): + """Reset this controller to default state. + + Will cancel any scheduled setup retry and will unload + the config entry. + """ + # If we have a retry scheduled, we were never setup. + if self._cancel_retry_setup is not None: + self._cancel_retry_setup() + self._cancel_retry_setup = None + return True + + # If the authentication was wrong. + if self.api is None: + return True + + if self.config_entry.data[CONF_POE_CONTROL]: + return await self.hass.config_entries.async_forward_entry_unload( + self.config_entry, 'switch') + return True + + +async def get_controller( + hass, host, username, password, port, site, verify_ssl): + """Create a controller object and verify authentication.""" + import aiounifi + + if verify_ssl: + session = aiohttp_client.async_get_clientsession(hass) + else: + session = aiohttp_client.async_create_clientsession( + hass, verify_ssl=verify_ssl, cookie_jar=CookieJar(unsafe=True)) + + controller = aiounifi.Controller( + host, username=username, password=password, port=port, site=site, + websession=session + ) + + try: + with async_timeout.timeout(5): + await controller.login() + return controller + + except aiounifi.Unauthorized: + LOGGER.warning("Connected to UniFi at %s but not registered.", host) + raise AuthenticationRequired + + except (asyncio.TimeoutError, aiounifi.RequestError): + LOGGER.error("Error connecting to the UniFi controller at %s", host) + raise CannotConnect + + except aiounifi.AiounifiException: + LOGGER.exception('Unknown UniFi communication error occurred') + raise AuthenticationRequired diff --git a/homeassistant/components/unifi/errors.py b/homeassistant/components/unifi/errors.py new file mode 100644 index 00000000000..c90c4956312 --- /dev/null +++ b/homeassistant/components/unifi/errors.py @@ -0,0 +1,26 @@ +"""Errors for the UniFi component.""" +from homeassistant.exceptions import HomeAssistantError + + +class UnifiException(HomeAssistantError): + """Base class for UniFi exceptions.""" + + +class AlreadyConfigured(UnifiException): + """Controller is already configured.""" + + +class AuthenticationRequired(UnifiException): + """Unknown error occurred.""" + + +class CannotConnect(UnifiException): + """Unable to connect to the controller.""" + + +class LoginRequired(UnifiException): + """Component got logged out.""" + + +class UserLevel(UnifiException): + """User level too low.""" diff --git a/homeassistant/components/unifi/strings.json b/homeassistant/components/unifi/strings.json new file mode 100644 index 00000000000..938ac058d22 --- /dev/null +++ b/homeassistant/components/unifi/strings.json @@ -0,0 +1,26 @@ +{ + "config": { + "title": "UniFi Controller", + "step": { + "user": { + "title": "Set up UniFi Controller", + "data": { + "host": "Host", + "username": "User name", + "password": "Password", + "port": "Port", + "site": "Site ID", + "verify_ssl": "Controller using proper certificate" + } + } + }, + "error": { + "faulty_credentials": "Bad user credentials", + "service_unavailable": "No service available" + }, + "abort": { + "already_configured": "Controller site is already configured", + "user_privilege": "User needs to be administrator" + } + } +} \ No newline at end of file diff --git a/homeassistant/config_entries.py b/homeassistant/config_entries.py index a4f28b63fb1..c1c0fbbf775 100644 --- a/homeassistant/config_entries.py +++ b/homeassistant/config_entries.py @@ -150,6 +150,7 @@ FLOWS = [ 'smhi', 'sonos', 'tradfri', + 'unifi', 'upnp', 'zone', 'zwave' diff --git a/requirements_all.txt b/requirements_all.txt index 3a488d16414..3c2d1df2545 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -118,6 +118,9 @@ aiolifx_effects==0.2.1 # homeassistant.components.scene.hunterdouglas_powerview aiopvapi==1.5.4 +# homeassistant.components.unifi +aiounifi==3 + # homeassistant.components.cover.aladdin_connect aladdin_connect==0.3 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index 51f6de977ef..9dd81e1a6c9 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -40,6 +40,9 @@ aiohttp_cors==0.7.0 # homeassistant.components.hue aiohue==1.5.0 +# homeassistant.components.unifi +aiounifi==3 + # homeassistant.components.notify.apns apns2==0.3.0 diff --git a/script/gen_requirements_all.py b/script/gen_requirements_all.py index fd8d673f633..47d11dff582 100755 --- a/script/gen_requirements_all.py +++ b/script/gen_requirements_all.py @@ -40,6 +40,7 @@ TEST_REQUIREMENTS = ( 'aioautomatic', 'aiohttp_cors', 'aiohue', + 'aiounifi', 'apns2', 'caldav', 'coinmarketcap', diff --git a/tests/components/switch/test_unifi.py b/tests/components/switch/test_unifi.py new file mode 100644 index 00000000000..f50bda34883 --- /dev/null +++ b/tests/components/switch/test_unifi.py @@ -0,0 +1,345 @@ +"""UniFi POE control platform tests.""" +from collections import deque +from unittest.mock import Mock + +import pytest + +import aiounifi +from aiounifi.clients import Clients +from aiounifi.devices import Devices + +from homeassistant import config_entries +from homeassistant.components import unifi +from homeassistant.setup import async_setup_component + +import homeassistant.components.switch as switch + +from tests.common import mock_coro + +CLIENT_1 = { + 'hostname': 'client_1', + 'ip': '10.0.0.1', + 'is_wired': True, + 'mac': '00:00:00:00:00:01', + 'name': 'POE Client 1', + 'oui': 'Producer', + 'sw_mac': '00:00:00:00:01:01', + 'sw_port': 1, + 'wired-rx_bytes': 1234000000, + 'wired-tx_bytes': 5678000000 +} +CLIENT_2 = { + 'hostname': 'client_2', + 'ip': '10.0.0.2', + 'is_wired': True, + 'mac': '00:00:00:00:00:02', + 'name': 'POE Client 2', + 'oui': 'Producer', + 'sw_mac': '00:00:00:00:01:01', + 'sw_port': 2, + 'wired-rx_bytes': 1234000000, + 'wired-tx_bytes': 5678000000 +} +CLIENT_3 = { + 'hostname': 'client_3', + 'ip': '10.0.0.3', + 'is_wired': True, + 'mac': '00:00:00:00:00:03', + 'name': 'Non-POE Client 3', + 'oui': 'Producer', + 'sw_mac': '00:00:00:00:01:01', + 'sw_port': 3, + 'wired-rx_bytes': 1234000000, + 'wired-tx_bytes': 5678000000 +} +CLIENT_4 = { + 'hostname': 'client_4', + 'ip': '10.0.0.4', + 'is_wired': True, + 'mac': '00:00:00:00:00:04', + 'name': 'Non-POE Client 4', + 'oui': 'Producer', + 'sw_mac': '00:00:00:00:01:01', + 'sw_port': 4, + 'wired-rx_bytes': 1234000000, + 'wired-tx_bytes': 5678000000 +} +POE_SWITCH_CLIENTS = [ + { + 'hostname': 'client_1', + 'ip': '10.0.0.1', + 'is_wired': True, + 'mac': '00:00:00:00:00:01', + 'name': 'POE Client 1', + 'oui': 'Producer', + 'sw_mac': '00:00:00:00:01:01', + 'sw_port': 1, + 'wired-rx_bytes': 1234000000, + 'wired-tx_bytes': 5678000000 + }, + { + 'hostname': 'client_2', + 'ip': '10.0.0.2', + 'is_wired': True, + 'mac': '00:00:00:00:00:02', + 'name': 'POE Client 2', + 'oui': 'Producer', + 'sw_mac': '00:00:00:00:01:01', + 'sw_port': 1, + 'wired-rx_bytes': 1234000000, + 'wired-tx_bytes': 5678000000 + } +] + +DEVICE_1 = { + 'device_id': 'mock-id', + 'ip': '10.0.1.1', + 'mac': '00:00:00:00:01:01', + 'type': 'usw', + 'name': 'mock-name', + 'portconf_id': '', + 'port_table': [ + { + 'media': 'GE', + 'name': 'Port 1', + 'port_idx': 1, + 'poe_class': 'Class 4', + 'poe_enable': True, + 'poe_mode': 'auto', + 'poe_power': '2.56', + 'poe_voltage': '53.40', + 'portconf_id': '1a1', + 'port_poe': True, + 'up': True + }, + { + 'media': 'GE', + 'name': 'Port 2', + 'port_idx': 2, + 'poe_class': 'Class 4', + 'poe_enable': True, + 'poe_mode': 'auto', + 'poe_power': '2.56', + 'poe_voltage': '53.40', + 'portconf_id': '1a2', + 'port_poe': True, + 'up': True + }, + { + 'media': 'GE', + 'name': 'Port 3', + 'port_idx': 3, + 'poe_class': 'Unknown', + 'poe_enable': False, + 'poe_mode': 'off', + 'poe_power': '0.00', + 'poe_voltage': '0.00', + 'portconf_id': '1a3', + 'port_poe': False, + 'up': True + }, + { + 'media': 'GE', + 'name': 'Port 4', + 'port_idx': 4, + 'poe_class': 'Unknown', + 'poe_enable': False, + 'poe_mode': 'auto', + 'poe_power': '0.00', + 'poe_voltage': '0.00', + 'portconf_id': '1a4', + 'port_poe': True, + 'up': True + } + ] +} + +CONTROLLER_DATA = { + unifi.CONF_HOST: 'mock-host', + unifi.CONF_USERNAME: 'mock-user', + unifi.CONF_PASSWORD: 'mock-pswd', + unifi.CONF_PORT: 1234, + unifi.CONF_SITE_ID: 'mock-site', + unifi.CONF_VERIFY_SSL: True +} + +ENTRY_CONFIG = { + unifi.CONF_CONTROLLER: CONTROLLER_DATA, + unifi.CONF_POE_CONTROL: True +} + +CONTROLLER_ID = unifi.CONTROLLER_ID.format(host='mock-host', site='mock-site') + + +@pytest.fixture +def mock_controller(hass): + """Mock a UniFi Controller.""" + controller = Mock( + available=True, + api=Mock(), + spec=unifi.UniFiController + ) + controller.mock_requests = [] + + controller.mock_client_responses = deque() + controller.mock_device_responses = deque() + + async def mock_request(method, path, **kwargs): + kwargs['method'] = method + kwargs['path'] = path + controller.mock_requests.append(kwargs) + if path == 's/{site}/stat/sta': + return controller.mock_client_responses.popleft() + if path == 's/{site}/stat/device': + return controller.mock_device_responses.popleft() + return None + + controller.api.clients = Clients({}, mock_request) + controller.api.devices = Devices({}, mock_request) + + return controller + + +async def setup_controller(hass, mock_controller): + """Load the UniFi switch platform with the provided controller.""" + hass.config.components.add(unifi.DOMAIN) + hass.data[unifi.DOMAIN] = {CONTROLLER_ID: mock_controller} + config_entry = config_entries.ConfigEntry( + 1, unifi.DOMAIN, 'Mock Title', ENTRY_CONFIG, 'test', + config_entries.CONN_CLASS_LOCAL_POLL) + await hass.config_entries.async_forward_entry_setup(config_entry, 'switch') + # To flush out the service call to update the group + await hass.async_block_till_done() + + +async def test_platform_manually_configured(hass): + """Test that we do not discover anything or try to set up a bridge.""" + assert await async_setup_component(hass, switch.DOMAIN, { + 'switch': { + 'platform': 'unifi' + } + }) is True + assert unifi.DOMAIN not in hass.data + + +async def test_no_clients(hass, mock_controller): + """Test the update_clients function when no clients are found.""" + mock_controller.mock_client_responses.append({}) + await setup_controller(hass, mock_controller) + assert len(mock_controller.mock_requests) == 2 + assert not hass.states.async_all() + + +async def test_switches(hass, mock_controller): + """Test the update_items function with some lights.""" + mock_controller.mock_client_responses.append([CLIENT_1, CLIENT_4]) + mock_controller.mock_device_responses.append([DEVICE_1]) + await setup_controller(hass, mock_controller) + assert len(mock_controller.mock_requests) == 2 + # 1 All Lights group, 2 lights + assert len(hass.states.async_all()) == 2 + + switch_1 = hass.states.get('switch.client_1') + assert switch_1 is not None + assert switch_1.state == 'on' + assert switch_1.attributes['power'] == '2.56' + assert switch_1.attributes['received'] == 1234 + assert switch_1.attributes['sent'] == 5678 + assert switch_1.attributes['switch'] == '00:00:00:00:01:01' + assert switch_1.attributes['port'] == 1 + assert switch_1.attributes['poe_mode'] == 'auto' + + switch = hass.states.get('switch.client_4') + assert switch is None + + +async def test_new_client_discovered(hass, mock_controller): + """Test if 2nd update has a new client.""" + mock_controller.mock_client_responses.append([CLIENT_1]) + mock_controller.mock_device_responses.append([DEVICE_1]) + + await setup_controller(hass, mock_controller) + assert len(mock_controller.mock_requests) == 2 + assert len(hass.states.async_all()) == 2 + + mock_controller.mock_client_responses.append([CLIENT_1, CLIENT_2]) + mock_controller.mock_device_responses.append([DEVICE_1]) + + # Calling a service will trigger the updates to run + await hass.services.async_call('switch', 'turn_off', { + 'entity_id': 'switch.client_1' + }, blocking=True) + # 2x light update, 1 turn on request + assert len(mock_controller.mock_requests) == 5 + assert len(hass.states.async_all()) == 3 + + switch = hass.states.get('switch.client_2') + assert switch is not None + assert switch.state == 'on' + + +async def test_failed_update_successful_login(hass, mock_controller): + """Running update can login when requested.""" + mock_controller.available = False + mock_controller.api.clients.update = Mock() + mock_controller.api.clients.update.side_effect = aiounifi.LoginRequired + mock_controller.api.login = Mock() + mock_controller.api.login.return_value = mock_coro() + + await setup_controller(hass, mock_controller) + assert len(mock_controller.mock_requests) == 0 + + assert mock_controller.available is True + + +async def test_failed_update_failed_login(hass, mock_controller): + """Running update can handle a failed login.""" + mock_controller.api.clients.update = Mock() + mock_controller.api.clients.update.side_effect = aiounifi.LoginRequired + mock_controller.api.login = Mock() + mock_controller.api.login.side_effect = aiounifi.AiounifiException + + await setup_controller(hass, mock_controller) + assert len(mock_controller.mock_requests) == 0 + + assert mock_controller.available is False + + +async def test_failed_update_unreachable_controller(hass, mock_controller): + """Running update can handle a unreachable controller.""" + mock_controller.mock_client_responses.append([CLIENT_1, CLIENT_2]) + mock_controller.mock_device_responses.append([DEVICE_1]) + + await setup_controller(hass, mock_controller) + + mock_controller.api.clients.update = Mock() + mock_controller.api.clients.update.side_effect = aiounifi.AiounifiException + + # Calling a service will trigger the updates to run + await hass.services.async_call('switch', 'turn_off', { + 'entity_id': 'switch.client_1' + }, blocking=True) + # 2x light update, 1 turn on request + assert len(mock_controller.mock_requests) == 3 + assert len(hass.states.async_all()) == 3 + + assert mock_controller.available is False + + +async def test_ignore_multiple_poe_clients_on_same_port(hass, mock_controller): + """Ignore when there are multiple POE driven clients on same port. + + If there is a non-UniFi switch powered by POE, + clients will be transparently marked as having POE as well. + """ + mock_controller.mock_client_responses.append(POE_SWITCH_CLIENTS) + mock_controller.mock_device_responses.append([DEVICE_1]) + await setup_controller(hass, mock_controller) + assert len(mock_controller.mock_requests) == 2 + # 1 All Lights group, 2 lights + assert len(hass.states.async_all()) == 0 + + switch_1 = hass.states.get('switch.client_1') + switch_2 = hass.states.get('switch.client_2') + assert switch_1 is None + assert switch_2 is None diff --git a/tests/components/unifi/__init__.py b/tests/components/unifi/__init__.py new file mode 100644 index 00000000000..e75b2778d2b --- /dev/null +++ b/tests/components/unifi/__init__.py @@ -0,0 +1 @@ +"""Tests for the UniFi component.""" diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py new file mode 100644 index 00000000000..b3b222d902a --- /dev/null +++ b/tests/components/unifi/test_controller.py @@ -0,0 +1,266 @@ +"""Test UniFi Controller.""" +from unittest.mock import Mock, patch + +from homeassistant.components import unifi +from homeassistant.components.unifi import controller, errors + +from tests.common import mock_coro + +CONTROLLER_DATA = { + unifi.CONF_HOST: '1.2.3.4', + unifi.CONF_USERNAME: 'username', + unifi.CONF_PASSWORD: 'password', + unifi.CONF_PORT: 1234, + unifi.CONF_SITE_ID: 'site', + unifi.CONF_VERIFY_SSL: True +} + +ENTRY_CONFIG = { + unifi.CONF_CONTROLLER: CONTROLLER_DATA, + unifi.CONF_POE_CONTROL: True + } + + +async def test_controller_setup(): + """Successful setup.""" + hass = Mock() + entry = Mock() + entry.data = ENTRY_CONFIG + api = Mock() + api.initialize.return_value = mock_coro(True) + + unifi_controller = controller.UniFiController(hass, entry) + + with patch.object(controller, 'get_controller', + return_value=mock_coro(api)): + assert await unifi_controller.async_setup() is True + + assert unifi_controller.api is api + assert len(hass.config_entries.async_forward_entry_setup.mock_calls) == 1 + assert hass.config_entries.async_forward_entry_setup.mock_calls[0][1] == \ + (entry, 'switch') + + +async def test_controller_host(): + """Config entry host and controller host are the same.""" + hass = Mock() + entry = Mock() + entry.data = ENTRY_CONFIG + + unifi_controller = controller.UniFiController(hass, entry) + + assert unifi_controller.host == '1.2.3.4' + + +async def test_controller_mac(): + """Test that it is possible to identify controller mac.""" + hass = Mock() + entry = Mock() + entry.data = ENTRY_CONFIG + client = Mock() + client.ip = '1.2.3.4' + client.mac = '00:11:22:33:44:55' + api = Mock() + api.initialize.return_value = mock_coro(True) + api.clients = {'client1': client} + + unifi_controller = controller.UniFiController(hass, entry) + + with patch.object(controller, 'get_controller', + return_value=mock_coro(api)): + assert await unifi_controller.async_setup() is True + + assert unifi_controller.mac == '00:11:22:33:44:55' + + +async def test_controller_no_mac(): + """Test that it works to not find the controllers mac.""" + hass = Mock() + entry = Mock() + entry.data = ENTRY_CONFIG + client = Mock() + client.ip = '5.6.7.8' + api = Mock() + api.initialize.return_value = mock_coro(True) + api.clients = {'client1': client} + + unifi_controller = controller.UniFiController(hass, entry) + + with patch.object(controller, 'get_controller', + return_value=mock_coro(api)): + assert await unifi_controller.async_setup() is True + + assert unifi_controller.mac is None + + +async def test_controller_not_accessible(): + """Retry to login gets scheduled when connection fails.""" + hass = Mock() + entry = Mock() + entry.data = ENTRY_CONFIG + api = Mock() + api.initialize.return_value = mock_coro(True) + + unifi_controller = controller.UniFiController(hass, entry) + + with patch.object(controller, 'get_controller', + side_effect=errors.CannotConnect): + assert await unifi_controller.async_setup() is False + + assert len(hass.helpers.event.async_call_later.mock_calls) == 1 + # Assert we are going to wait 2 seconds + assert hass.helpers.event.async_call_later.mock_calls[0][1][0] == 2 + + +async def test_controller_unknown_error(): + """Unknown errors are handled.""" + hass = Mock() + entry = Mock() + entry.data = ENTRY_CONFIG + api = Mock() + api.initialize.return_value = mock_coro(True) + + unifi_controller = controller.UniFiController(hass, entry) + + with patch.object(controller, 'get_controller', side_effect=Exception): + assert await unifi_controller.async_setup() is False + + assert not hass.helpers.event.async_call_later.mock_calls + + +async def test_reset_cancels_retry_setup(): + """Resetting a controller while we're waiting to retry setup.""" + hass = Mock() + entry = Mock() + entry.data = ENTRY_CONFIG + + unifi_controller = controller.UniFiController(hass, entry) + + with patch.object(controller, 'get_controller', + side_effect=errors.CannotConnect): + assert await unifi_controller.async_setup() is False + + mock_call_later = hass.helpers.event.async_call_later + + assert len(mock_call_later.mock_calls) == 1 + + assert await unifi_controller.async_reset() + + assert len(mock_call_later.mock_calls) == 2 + assert len(mock_call_later.return_value.mock_calls) == 1 + + +async def test_reset_if_entry_had_wrong_auth(): + """Calling reset when the entry contains wrong auth.""" + hass = Mock() + entry = Mock() + entry.data = ENTRY_CONFIG + + unifi_controller = controller.UniFiController(hass, entry) + + with patch.object(controller, 'get_controller', + side_effect=errors.AuthenticationRequired): + assert await unifi_controller.async_setup() is False + + assert not hass.async_add_job.mock_calls + + assert await unifi_controller.async_reset() + + +async def test_reset_unloads_entry_if_setup(): + """Calling reset when the entry has been setup.""" + hass = Mock() + entry = Mock() + entry.data = ENTRY_CONFIG + api = Mock() + api.initialize.return_value = mock_coro(True) + + unifi_controller = controller.UniFiController(hass, entry) + + with patch.object(controller, 'get_controller', + return_value=mock_coro(api)): + assert await unifi_controller.async_setup() is True + + assert len(hass.config_entries.async_forward_entry_setup.mock_calls) == 1 + + hass.config_entries.async_forward_entry_unload.return_value = \ + mock_coro(True) + assert await unifi_controller.async_reset() + + assert len(hass.config_entries.async_forward_entry_unload.mock_calls) == 1 + + +async def test_reset_unloads_entry_without_poe_control(): + """Calling reset while the entry has been setup.""" + hass = Mock() + entry = Mock() + entry.data = dict(ENTRY_CONFIG) + entry.data[unifi.CONF_POE_CONTROL] = False + api = Mock() + api.initialize.return_value = mock_coro(True) + + unifi_controller = controller.UniFiController(hass, entry) + + with patch.object(controller, 'get_controller', + return_value=mock_coro(api)): + assert await unifi_controller.async_setup() is True + + assert not hass.config_entries.async_forward_entry_setup.mock_calls + + hass.config_entries.async_forward_entry_unload.return_value = \ + mock_coro(True) + assert await unifi_controller.async_reset() + + assert not hass.config_entries.async_forward_entry_unload.mock_calls + + +async def test_get_controller(hass): + """Successful call.""" + with patch('aiounifi.Controller.login', return_value=mock_coro()): + assert await controller.get_controller(hass, **CONTROLLER_DATA) + + +async def test_get_controller_verify_ssl_false(hass): + """Successful call with verify ssl set to false.""" + controller_data = dict(CONTROLLER_DATA) + controller_data[unifi.CONF_VERIFY_SSL] = False + with patch('aiounifi.Controller.login', return_value=mock_coro()): + assert await controller.get_controller(hass, **controller_data) + + +async def test_get_controller_login_failed(hass): + """Check that get_controller can handle a failed login.""" + import aiounifi + result = None + with patch('aiounifi.Controller.login', side_effect=aiounifi.Unauthorized): + try: + result = await controller.get_controller(hass, **CONTROLLER_DATA) + except errors.AuthenticationRequired: + pass + assert result is None + + +async def test_get_controller_controller_unavailable(hass): + """Check that get_controller can handle controller being unavailable.""" + import aiounifi + result = None + with patch('aiounifi.Controller.login', + side_effect=aiounifi.RequestError): + try: + result = await controller.get_controller(hass, **CONTROLLER_DATA) + except errors.CannotConnect: + pass + assert result is None + + +async def test_get_controller_unknown_error(hass): + """Check that get_controller can handle unkown errors.""" + import aiounifi + result = None + with patch('aiounifi.Controller.login', + side_effect=aiounifi.AiounifiException): + try: + result = await controller.get_controller(hass, **CONTROLLER_DATA) + except errors.AuthenticationRequired: + pass + assert result is None diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py new file mode 100644 index 00000000000..400dd3fd93e --- /dev/null +++ b/tests/components/unifi/test_init.py @@ -0,0 +1,330 @@ +"""Test UniFi setup process.""" +from unittest.mock import Mock, patch + +from homeassistant.components import unifi +from homeassistant.setup import async_setup_component + +from tests.common import mock_coro, MockConfigEntry + + +async def test_setup_with_no_config(hass): + """Test that we do not discover anything or try to set up a bridge.""" + assert await async_setup_component(hass, unifi.DOMAIN, {}) is True + assert unifi.DOMAIN not in hass.data + + +async def test_successful_config_entry(hass): + """Test that configured options for a host are loaded via config entry.""" + entry = MockConfigEntry(domain=unifi.DOMAIN, data={ + 'controller': { + 'host': '0.0.0.0', + 'username': 'user', + 'password': 'pass', + 'port': 80, + 'site': 'default', + 'verify_ssl': True + }, + 'poe_control': True + }) + entry.add_to_hass(hass) + mock_registry = Mock() + with patch.object(unifi, 'UniFiController') as mock_controller, \ + patch('homeassistant.helpers.device_registry.async_get_registry', + return_value=mock_coro(mock_registry)): + mock_controller.return_value.async_setup.return_value = mock_coro(True) + mock_controller.return_value.mac = '00:11:22:33:44:55' + assert await unifi.async_setup_entry(hass, entry) is True + + assert len(mock_controller.mock_calls) == 2 + p_hass, p_entry = mock_controller.mock_calls[0][1] + + assert p_hass is hass + assert p_entry is entry + + assert len(mock_registry.mock_calls) == 1 + assert mock_registry.mock_calls[0][2] == { + 'config_entry_id': entry.entry_id, + 'connections': { + ('mac', '00:11:22:33:44:55') + }, + 'manufacturer': 'Ubiquiti', + 'model': "UniFi Controller", + 'name': "UniFi Controller", + } + + +async def test_controller_fail_setup(hass): + """Test that configured options for a host are loaded via config entry.""" + entry = MockConfigEntry(domain=unifi.DOMAIN, data={ + 'controller': { + 'host': '0.0.0.0', + 'username': 'user', + 'password': 'pass', + 'port': 80, + 'site': 'default', + 'verify_ssl': True + }, + 'poe_control': True + }) + entry.add_to_hass(hass) + + with patch.object(unifi, 'UniFiController') as mock_cntrlr: + mock_cntrlr.return_value.async_setup.return_value = mock_coro(False) + assert await unifi.async_setup_entry(hass, entry) is False + + controller_id = unifi.CONTROLLER_ID.format( + host='0.0.0.0', site='default' + ) + assert controller_id not in hass.data[unifi.DOMAIN] + + +async def test_controller_no_mac(hass): + """Test that configured options for a host are loaded via config entry.""" + entry = MockConfigEntry(domain=unifi.DOMAIN, data={ + 'controller': { + 'host': '0.0.0.0', + 'username': 'user', + 'password': 'pass', + 'port': 80, + 'site': 'default', + 'verify_ssl': True + }, + 'poe_control': True + }) + entry.add_to_hass(hass) + mock_registry = Mock() + with patch.object(unifi, 'UniFiController') as mock_controller, \ + patch('homeassistant.helpers.device_registry.async_get_registry', + return_value=mock_coro(mock_registry)): + mock_controller.return_value.async_setup.return_value = mock_coro(True) + mock_controller.return_value.mac = None + assert await unifi.async_setup_entry(hass, entry) is True + + assert len(mock_controller.mock_calls) == 2 + + assert len(mock_registry.mock_calls) == 0 + + +async def test_unload_entry(hass): + """Test being able to unload an entry.""" + entry = MockConfigEntry(domain=unifi.DOMAIN, data={ + 'controller': { + 'host': '0.0.0.0', + 'username': 'user', + 'password': 'pass', + 'port': 80, + 'site': 'default', + 'verify_ssl': True + }, + 'poe_control': True + }) + entry.add_to_hass(hass) + + with patch.object(unifi, 'UniFiController') as mock_controller, \ + patch('homeassistant.helpers.device_registry.async_get_registry', + return_value=mock_coro(Mock())): + mock_controller.return_value.async_setup.return_value = mock_coro(True) + mock_controller.return_value.mac = '00:11:22:33:44:55' + assert await unifi.async_setup_entry(hass, entry) is True + + assert len(mock_controller.return_value.mock_calls) == 1 + + mock_controller.return_value.async_reset.return_value = mock_coro(True) + assert await unifi.async_unload_entry(hass, entry) + assert len(mock_controller.return_value.async_reset.mock_calls) == 1 + assert hass.data[unifi.DOMAIN] == {} + + +async def test_flow_works(hass, aioclient_mock): + """Test config flow.""" + flow = unifi.UnifiFlowHandler() + flow.hass = hass + + with patch('aiounifi.Controller') as mock_controller: + def mock_constructor(host, username, password, port, site, websession): + """Fake the controller constructor.""" + mock_controller.host = host + mock_controller.username = username + mock_controller.password = password + mock_controller.port = port + mock_controller.site = site + return mock_controller + + mock_controller.side_effect = mock_constructor + mock_controller.login.return_value = mock_coro() + mock_controller.sites.return_value = mock_coro({ + 'site1': {'name': 'default', 'role': 'admin', 'desc': 'site name'} + }) + + await flow.async_step_user(user_input={ + unifi.CONF_HOST: '1.2.3.4', + unifi.CONF_USERNAME: 'username', + unifi.CONF_PASSWORD: 'password', + unifi.CONF_PORT: 1234, + unifi.CONF_VERIFY_SSL: True + }) + + result = await flow.async_step_site(user_input={}) + + assert mock_controller.host == '1.2.3.4' + assert len(mock_controller.login.mock_calls) == 1 + assert len(mock_controller.sites.mock_calls) == 1 + + assert result['type'] == 'create_entry' + assert result['title'] == 'site name' + assert result['data'] == { + unifi.CONF_CONTROLLER: { + unifi.CONF_HOST: '1.2.3.4', + unifi.CONF_USERNAME: 'username', + unifi.CONF_PASSWORD: 'password', + unifi.CONF_PORT: 1234, + unifi.CONF_SITE_ID: 'default', + unifi.CONF_VERIFY_SSL: True + }, + unifi.CONF_POE_CONTROL: True + } + + +async def test_controller_multiple_sites(hass): + """Test config flow.""" + flow = unifi.UnifiFlowHandler() + flow.hass = hass + + flow.config = { + unifi.CONF_HOST: '1.2.3.4', + unifi.CONF_USERNAME: 'username', + unifi.CONF_PASSWORD: 'password', + } + flow.sites = { + 'site1': { + 'name': 'default', 'role': 'admin', 'desc': 'site name' + }, + 'site2': { + 'name': 'site2', 'role': 'admin', 'desc': 'site2 name' + } + } + + result = await flow.async_step_site() + + assert result['type'] == 'form' + assert result['step_id'] == 'site' + + assert result['data_schema']({'site': 'site name'}) + assert result['data_schema']({'site': 'site2 name'}) + + +async def test_controller_site_already_configured(hass): + """Test config flow.""" + flow = unifi.UnifiFlowHandler() + flow.hass = hass + + entry = MockConfigEntry(domain=unifi.DOMAIN, data={ + 'controller': { + 'host': '1.2.3.4', + 'site': 'default', + } + }) + entry.add_to_hass(hass) + + flow.config = { + unifi.CONF_HOST: '1.2.3.4', + unifi.CONF_USERNAME: 'username', + unifi.CONF_PASSWORD: 'password', + } + flow.desc = 'site name' + flow.sites = { + 'site1': { + 'name': 'default', 'role': 'admin', 'desc': 'site name' + } + } + + result = await flow.async_step_site() + + assert result['type'] == 'abort' + + +async def test_user_permissions_low(hass, aioclient_mock): + """Test config flow.""" + flow = unifi.UnifiFlowHandler() + flow.hass = hass + + with patch('aiounifi.Controller') as mock_controller: + def mock_constructor(host, username, password, port, site, websession): + """Fake the controller constructor.""" + mock_controller.host = host + mock_controller.username = username + mock_controller.password = password + mock_controller.port = port + mock_controller.site = site + return mock_controller + + mock_controller.side_effect = mock_constructor + mock_controller.login.return_value = mock_coro() + mock_controller.sites.return_value = mock_coro({ + 'site1': {'name': 'default', 'role': 'viewer', 'desc': 'site name'} + }) + + await flow.async_step_user(user_input={ + unifi.CONF_HOST: '1.2.3.4', + unifi.CONF_USERNAME: 'username', + unifi.CONF_PASSWORD: 'password', + unifi.CONF_PORT: 1234, + unifi.CONF_VERIFY_SSL: True + }) + + result = await flow.async_step_site(user_input={}) + + assert result['type'] == 'abort' + + +async def test_user_credentials_faulty(hass, aioclient_mock): + """Test config flow.""" + flow = unifi.UnifiFlowHandler() + flow.hass = hass + + with patch.object(unifi, 'get_controller', + side_effect=unifi.errors.AuthenticationRequired): + result = await flow.async_step_user({ + unifi.CONF_HOST: '1.2.3.4', + unifi.CONF_USERNAME: 'username', + unifi.CONF_PASSWORD: 'password', + unifi.CONF_SITE_ID: 'default', + }) + + assert result['type'] == 'form' + assert result['errors'] == {'base': 'faulty_credentials'} + + +async def test_controller_is_unavailable(hass, aioclient_mock): + """Test config flow.""" + flow = unifi.UnifiFlowHandler() + flow.hass = hass + + with patch.object(unifi, 'get_controller', + side_effect=unifi.errors.CannotConnect): + result = await flow.async_step_user({ + unifi.CONF_HOST: '1.2.3.4', + unifi.CONF_USERNAME: 'username', + unifi.CONF_PASSWORD: 'password', + unifi.CONF_SITE_ID: 'default', + }) + + assert result['type'] == 'form' + assert result['errors'] == {'base': 'service_unavailable'} + + +async def test_controller_unkown_problem(hass, aioclient_mock): + """Test config flow.""" + flow = unifi.UnifiFlowHandler() + flow.hass = hass + + with patch.object(unifi, 'get_controller', + side_effect=Exception): + result = await flow.async_step_user({ + unifi.CONF_HOST: '1.2.3.4', + unifi.CONF_USERNAME: 'username', + unifi.CONF_PASSWORD: 'password', + unifi.CONF_SITE_ID: 'default', + }) + + assert result['type'] == 'abort'