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

@@ -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
 

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]

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"
+        }
+    }
+}

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,
+        )

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'

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

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."""

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"
+        }
+    }
+}

homeassistant/config_entries.py

@@ -150,6 +150,7 @@ FLOWS = [
     'smhi',
     'sonos',
     'tradfri',
+    'unifi',
     'upnp',
     'zone',
     'zwave'

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
 

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
 

script/gen_requirements_all.py

@@ -40,6 +40,7 @@ TEST_REQUIREMENTS = (
     'aioautomatic',
     'aiohttp_cors',
     'aiohue',
+    'aiounifi',
     'apns2',
     'caldav',
     'coinmarketcap',

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

tests/components/unifi/__init__.py

@@ -0,0 +1 @@
+"""Tests for the UniFi component."""

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

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'