Merge UniFi device tracker to config entry (#24367)

* Move device tracker to use config entry

* Remove monitored conditions attributes based on ADR0003

* Add support for import of device tracker config to be backwards compatible

* Remove unnecessary configuration options from device tracker

* Add component configuration support

homeassistant/components/device_tracker/config_entry.py

@@ -1,8 +1,7 @@
 """Code to set up a device tracker platform using a config entry."""
 from typing import Optional
 
-from homeassistant.helpers.entity import Entity
+from homeassistant.components import zone
-from homeassistant.helpers.entity_component import EntityComponent
 from homeassistant.const import (
     STATE_NOT_HOME,
     STATE_HOME,
@@ -11,7 +10,8 @@ from homeassistant.const import (
     ATTR_LONGITUDE,
     ATTR_BATTERY_LEVEL,
 )
-from homeassistant.components import zone
+from homeassistant.helpers.entity import Entity
+from homeassistant.helpers.entity_component import EntityComponent
 
 from .const import (
     ATTR_SOURCE_TYPE,

homeassistant/components/unifi/__init__.py

@@ -1,13 +1,41 @@
 """Support for devices connected to UniFi POE."""
+import voluptuous as vol
+
 from homeassistant.const import CONF_HOST
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 
-from .const import CONF_CONTROLLER, CONF_SITE_ID, CONTROLLER_ID, DOMAIN
+import homeassistant.helpers.config_validation as cv
+
+from .const import (
+    CONF_CONTROLLER, CONF_DETECTION_TIME, CONF_SITE_ID, CONF_SSID_FILTER,
+    CONTROLLER_ID, DOMAIN, UNIFI_CONFIG)
 from .controller import UniFiController
 
+CONF_CONTROLLERS = 'controllers'
+
+CONTROLLER_SCHEMA = vol.Schema({
+    vol.Required(CONF_HOST): cv.string,
+    vol.Required(CONF_SITE_ID): cv.string,
+    vol.Optional(CONF_DETECTION_TIME): vol.All(
+        cv.time_period, cv.positive_timedelta),
+    vol.Optional(CONF_SSID_FILTER): vol.All(cv.ensure_list, [cv.string])
+})
+
+CONFIG_SCHEMA = vol.Schema({
+    DOMAIN: vol.Schema({
+        vol.Required(CONF_CONTROLLERS):
+            vol.All(cv.ensure_list, [CONTROLLER_SCHEMA]),
+    }),
+}, extra=vol.ALLOW_EXTRA)
+
 
 async def async_setup(hass, config):
     """Component doesn't support configuration through configuration.yaml."""
+    hass.data[UNIFI_CONFIG] = []
+
+    if DOMAIN in config:
+        hass.data[UNIFI_CONFIG] = config[DOMAIN][CONF_CONTROLLERS]
+
     return True
 
 

homeassistant/components/unifi/config_flow.py

@@ -7,9 +7,7 @@ from homeassistant.const import (
 
 from .const import CONF_CONTROLLER, CONF_SITE_ID, DOMAIN, LOGGER
 from .controller import get_controller
-from .errors import (
+from .errors import AlreadyConfigured, AuthenticationRequired, CannotConnect
-    AlreadyConfigured, AuthenticationRequired, CannotConnect, UserLevel)
-
 
 DEFAULT_PORT = 8443
 DEFAULT_SITE_ID = 'default'
@@ -44,6 +42,7 @@ class UnifiFlowHandler(config_entries.ConfigFlow):
                     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()
@@ -80,14 +79,11 @@ class UnifiFlowHandler(config_entries.ConfigFlow):
         errors = {}
 
         if user_input is not None:
-
             try:
                 desc = user_input.get(CONF_SITE_ID, self.desc)
-                print(self.sites)
+
                 for site in self.sites.values():
                     if desc == site['desc']:
-                        if site['role'] != 'admin':
-                            raise UserLevel
                         self.config[CONF_SITE_ID] = site['name']
                         break
 
@@ -109,13 +105,16 @@ class UnifiFlowHandler(config_entries.ConfigFlow):
             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={})
 
+        if self.desc is not None:
+            for site in self.sites.values():
+                if self.desc == site['name']:
+                    self.desc = site['desc']
+                    return await self.async_step_site(user_input={})
+
         sites = []
         for site in self.sites.values():
             sites.append(site['desc'])
@@ -127,3 +126,17 @@ class UnifiFlowHandler(config_entries.ConfigFlow):
             }),
             errors=errors,
         )
+
+    async def async_step_import(self, import_config):
+        """Import from UniFi device tracker config."""
+        config = {
+            CONF_HOST: import_config[CONF_HOST],
+            CONF_USERNAME: import_config[CONF_USERNAME],
+            CONF_PASSWORD: import_config[CONF_PASSWORD],
+            CONF_PORT: import_config.get(CONF_PORT),
+            CONF_VERIFY_SSL: import_config.get(CONF_VERIFY_SSL),
+        }
+
+        self.desc = import_config[CONF_SITE_ID]
+
+        return await self.async_step_user(user_input=config)

homeassistant/components/unifi/const.py

@@ -8,3 +8,8 @@ CONTROLLER_ID = '{host}-{site}'
 
 CONF_CONTROLLER = 'controller'
 CONF_SITE_ID = 'site'
+
+UNIFI_CONFIG = 'unifi_config'
+
+CONF_DETECTION_TIME = 'detection_time'
+CONF_SSID_FILTER = 'ssid_filter'

homeassistant/components/unifi/controller.py

@@ -12,7 +12,8 @@ from homeassistant.const import CONF_HOST
 from homeassistant.helpers import aiohttp_client
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 
-from .const import CONF_CONTROLLER, CONF_SITE_ID, CONTROLLER_ID, LOGGER
+from .const import (
+    CONF_CONTROLLER, CONF_SITE_ID, CONTROLLER_ID, LOGGER, UNIFI_CONFIG)
 from .errors import AuthenticationRequired, CannotConnect
 
 
@@ -27,11 +28,30 @@ class UniFiController:
         self.api = None
         self.progress = None
 
+        self._site_name = None
+        self._site_role = None
+        self.unifi_config = {}
+
     @property
     def host(self):
         """Return the host of this controller."""
         return self.config_entry.data[CONF_CONTROLLER][CONF_HOST]
 
+    @property
+    def site(self):
+        """Return the site of this config entry."""
+        return self.config_entry.data[CONF_CONTROLLER][CONF_SITE_ID]
+
+    @property
+    def site_name(self):
+        """Return the nice name of site."""
+        return self._site_name
+
+    @property
+    def site_role(self):
+        """Return the site user role of this controller."""
+        return self._site_role
+
     @property
     def mac(self):
         """Return the mac address of this controller."""
@@ -44,9 +64,7 @@ class UniFiController:
     def event_update(self):
         """Event specific per UniFi entry to signal new data."""
         return 'unifi-update-{}'.format(
-            CONTROLLER_ID.format(
+            CONTROLLER_ID.format(host=self.host, site=self.site))
-                host=self.host,
-                site=self.config_entry.data[CONF_CONTROLLER][CONF_SITE_ID]))
 
     async def request_update(self):
         """Request an update."""
@@ -63,7 +81,7 @@ class UniFiController:
         failed = False
 
         try:
-            with async_timeout.timeout(4):
+            with async_timeout.timeout(10):
                 await self.api.clients.update()
                 await self.api.devices.update()
 
@@ -99,6 +117,14 @@ class UniFiController:
                 self.hass, **self.config_entry.data[CONF_CONTROLLER])
             await self.api.initialize()
 
+            sites = await self.api.sites()
+
+            for site in sites.values():
+                if self.site == site['name']:
+                    self._site_name = site['desc']
+                    self._site_role = site['role']
+                    break
+
         except CannotConnect:
             raise ConfigEntryNotReady
 
@@ -107,9 +133,16 @@ class UniFiController:
                 'Unknown error connecting with UniFi controller.')
             return False
 
-        hass.async_create_task(
+        for unifi_config in hass.data[UNIFI_CONFIG]:
-            hass.config_entries.async_forward_entry_setup(
+            if self.host == unifi_config[CONF_HOST] and \
-                self.config_entry, 'switch'))
+                    self.site == unifi_config[CONF_SITE_ID]:
+                self.unifi_config = unifi_config
+                break
+
+        for platform in ['device_tracker', 'switch']:
+            hass.async_create_task(
+                hass.config_entries.async_forward_entry_setup(
+                    self.config_entry, platform))
 
         return True
 
@@ -123,8 +156,11 @@ class UniFiController:
         if self.api is None:
             return True
 
-        return await self.hass.config_entries.async_forward_entry_unload(
+        for platform in ['device_tracker', 'switch']:
-            self.config_entry, 'switch')
+            await self.hass.config_entries.async_forward_entry_unload(
+                self.config_entry, platform)
+
+        return True
 
 
 async def get_controller(

homeassistant/components/unifi/device_tracker.py

@@ -1,172 +1,181 @@
 """Support for Unifi WAP controllers."""
-import asyncio
-import logging
 from datetime import timedelta
+
+import logging
 import voluptuous as vol
 
-import async_timeout
+from homeassistant import config_entries
-
+from homeassistant.components import unifi
-import aiounifi
+from homeassistant.components.device_tracker import PLATFORM_SCHEMA
+from homeassistant.components.device_tracker.config_entry import ScannerEntity
+from homeassistant.components.device_tracker.const import SOURCE_TYPE_ROUTER
+from homeassistant.core import callback
+from homeassistant.const import (
+    CONF_HOST, CONF_USERNAME, CONF_PASSWORD, CONF_PORT, CONF_VERIFY_SSL)
+from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
 
 import homeassistant.helpers.config_validation as cv
-from homeassistant.components.device_tracker import (
-    DOMAIN, PLATFORM_SCHEMA, DeviceScanner)
-from homeassistant.const import CONF_HOST, CONF_USERNAME, CONF_PASSWORD
-from homeassistant.const import CONF_VERIFY_SSL, CONF_MONITORED_CONDITIONS
 import homeassistant.util.dt as dt_util
 
-from .controller import get_controller
+from .const import (
-from .errors import AuthenticationRequired, CannotConnect
+    CONF_CONTROLLER, CONF_DETECTION_TIME, CONF_SITE_ID, CONF_SSID_FILTER,
+    CONTROLLER_ID, DOMAIN as UNIFI_DOMAIN)
 
-_LOGGER = logging.getLogger(__name__)
+LOGGER = logging.getLogger(__name__)
-CONF_PORT = 'port'
+
-CONF_SITE_ID = 'site_id'
+DEVICE_ATTRIBUTES = [
-CONF_DETECTION_TIME = 'detection_time'
+    '_is_guest_by_uap', 'ap_mac', 'authorized', 'bssid', 'ccq',
-CONF_SSID_FILTER = 'ssid_filter'
+    'channel', 'essid', 'hostname', 'ip', 'is_11r', 'is_guest', 'is_wired',
+    'mac', 'name', 'noise', 'noted', 'oui', 'qos_policy_applied', 'radio',
+    'radio_proto', 'rssi', 'signal', 'site_id', 'vlan'
+]
+
+CONF_DT_SITE_ID = 'site_id'
 
 DEFAULT_HOST = 'localhost'
 DEFAULT_PORT = 8443
 DEFAULT_VERIFY_SSL = True
 DEFAULT_DETECTION_TIME = timedelta(seconds=300)
 
-NOTIFICATION_ID = 'unifi_notification'
-NOTIFICATION_TITLE = 'Unifi Device Tracker Setup'
-
-AVAILABLE_ATTRS = [
-    '_id', '_is_guest_by_uap', '_last_seen_by_uap', '_uptime_by_uap',
-    'ap_mac', 'assoc_time', 'authorized', 'bssid', 'bytes-r', 'ccq',
-    'channel', 'essid', 'first_seen', 'hostname', 'idletime', 'ip',
-    'is_11r', 'is_guest', 'is_wired', 'last_seen', 'latest_assoc_time',
-    'mac', 'name', 'noise', 'noted', 'oui', 'powersave_enabled',
-    'qos_policy_applied', 'radio', 'radio_proto', 'rssi', 'rx_bytes',
-    'rx_bytes-r', 'rx_packets', 'rx_rate', 'signal', 'site_id',
-    'tx_bytes', 'tx_bytes-r', 'tx_packets', 'tx_power', 'tx_rate',
-    'uptime', 'user_id', 'usergroup_id', 'vlan'
-]
-
-TIMESTAMP_ATTRS = ['first_seen', 'last_seen', 'latest_assoc_time']
-
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
     vol.Optional(CONF_HOST, default=DEFAULT_HOST): cv.string,
-    vol.Optional(CONF_SITE_ID, default='default'): cv.string,
+    vol.Optional(CONF_DT_SITE_ID, default='default'): cv.string,
     vol.Required(CONF_PASSWORD): cv.string,
     vol.Required(CONF_USERNAME): cv.string,
     vol.Required(CONF_PORT, default=DEFAULT_PORT): cv.port,
     vol.Optional(CONF_VERIFY_SSL, default=DEFAULT_VERIFY_SSL): vol.Any(
-        cv.boolean, cv.isfile),
+        cv.boolean, cv.isfile)
-    vol.Optional(CONF_DETECTION_TIME, default=DEFAULT_DETECTION_TIME): vol.All(
+}, extra=vol.ALLOW_EXTRA)
-        cv.time_period, cv.positive_timedelta),
-    vol.Optional(CONF_MONITORED_CONDITIONS):
-        vol.All(cv.ensure_list, [vol.In(AVAILABLE_ATTRS)]),
-    vol.Optional(CONF_SSID_FILTER): vol.All(cv.ensure_list, [cv.string])
-})
 
 
-async def async_get_scanner(hass, config):
+async def async_setup_scanner(hass, config, sync_see, discovery_info):
-    """Set up the Unifi device_tracker."""
+    """Set up the Unifi integration."""
-    host = config[DOMAIN].get(CONF_HOST)
+    config[CONF_SITE_ID] = config.pop(CONF_DT_SITE_ID)  # Current from legacy
-    username = config[DOMAIN].get(CONF_USERNAME)
-    password = config[DOMAIN].get(CONF_PASSWORD)
-    site_id = config[DOMAIN].get(CONF_SITE_ID)
-    port = config[DOMAIN].get(CONF_PORT)
-    verify_ssl = config[DOMAIN].get(CONF_VERIFY_SSL)
-    detection_time = config[DOMAIN].get(CONF_DETECTION_TIME)
-    monitored_conditions = config[DOMAIN].get(CONF_MONITORED_CONDITIONS)
-    ssid_filter = config[DOMAIN].get(CONF_SSID_FILTER)
 
-    try:
+    exist = False
-        controller = await get_controller(
-            hass, host, username, password, port, site_id, verify_ssl)
-        await controller.initialize()
 
-    except (AuthenticationRequired, CannotConnect) as ex:
+    for entry in hass.config_entries.async_entries(UNIFI_DOMAIN):
-        _LOGGER.error("Failed to connect to Unifi: %s", ex)
+        if config[CONF_HOST] == entry.data[CONF_CONTROLLER][CONF_HOST] and \
-        hass.components.persistent_notification.create(
+                config[CONF_SITE_ID] == \
-            'Failed to connect to Unifi. '
+                entry.data[CONF_CONTROLLER][CONF_SITE_ID]:
-            'Error: {}<br />'
+            exist = True
-            'You will need to restart hass after fixing.'
+            break
-            ''.format(ex),
-            title=NOTIFICATION_TITLE,
-            notification_id=NOTIFICATION_ID)
-        return False
 
-    return UnifiScanner(
+    if not exist:
-        controller, detection_time, ssid_filter, monitored_conditions)
+        hass.async_create_task(hass.config_entries.flow.async_init(
+            UNIFI_DOMAIN, context={'source': config_entries.SOURCE_IMPORT},
+            data=config
+        ))
+
+    return True
 
 
-class UnifiScanner(DeviceScanner):
+async def async_setup_entry(hass, config_entry, async_add_entities):
-    """Provide device_tracker support from Unifi WAP client data."""
+    """Set up device tracker for UniFi component."""
+    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]
+    tracked = {}
 
-    def __init__(self, controller, detection_time: timedelta,
+    @callback
-                 ssid_filter, monitored_conditions) -> None:
+    def update_controller():
-        """Initialize the scanner."""
+        """Update the values of the controller."""
+        update_items(controller, async_add_entities, tracked)
+
+    async_dispatcher_connect(hass, controller.event_update, update_controller)
+
+    update_controller()
+
+
+@callback
+def update_items(controller, async_add_entities, tracked):
+    """Update tracked device state from the controller."""
+    new_tracked = []
+
+    for client_id in controller.api.clients:
+
+        if client_id in tracked:
+            LOGGER.debug("Updating UniFi tracked device %s (%s)",
+                         tracked[client_id].entity_id,
+                         tracked[client_id].client.mac)
+            tracked[client_id].async_schedule_update_ha_state()
+            continue
+
+        client = controller.api.clients[client_id]
+
+        if not client.is_wired and \
+                CONF_SSID_FILTER in controller.unifi_config and \
+                client.essid not in controller.unifi_config[CONF_SSID_FILTER]:
+            continue
+
+        tracked[client_id] = UniFiClientTracker(client, controller)
+        new_tracked.append(tracked[client_id])
+        LOGGER.debug("New UniFi switch %s (%s)", client.hostname, client.mac)
+
+    if new_tracked:
+        async_add_entities(new_tracked)
+
+
+class UniFiClientTracker(ScannerEntity):
+    """Representation of a network device."""
+
+    def __init__(self, client, controller):
+        """Set up tracked device."""
+        self.client = client
         self.controller = controller
-        self._detection_time = detection_time
-        self._ssid_filter = ssid_filter
-        self._monitored_conditions = monitored_conditions
-        self._clients = {}
 
     async def async_update(self):
-        """Get the clients from the device."""
+        """Synchronize state with controller."""
-        try:
+        await self.controller.request_update()
-            await self.controller.clients.update()
-            clients = self.controller.clients.values()
 
-        except aiounifi.LoginRequired:
+    @property
-            try:
+    def is_connected(self):
-                with async_timeout.timeout(5):
+        """Return true if the device is connected to the network."""
-                    await self.controller.login()
+        detection_time = self.controller.unifi_config.get(
-            except (asyncio.TimeoutError, aiounifi.AiounifiException):
+            CONF_DETECTION_TIME, DEFAULT_DETECTION_TIME)
-                clients = []
 
-        except aiounifi.AiounifiException:
+        if (dt_util.utcnow() - dt_util.utc_from_timestamp(float(
-            clients = []
+                self.client.last_seen))) < detection_time:
+            return True
+        return False
 
-        # Filter clients to provided SSID list
+    @property
-        if self._ssid_filter:
+    def source_type(self):
-            clients = [
+        """Return the source type of the device."""
-                client for client in clients
+        return SOURCE_TYPE_ROUTER
-                if client.essid in self._ssid_filter
-            ]
 
-        self._clients = {
+    @property
-            client.raw['mac']: client.raw
+    def name(self) -> str:
-            for client in clients
+        """Return the name of the device."""
-            if (dt_util.utcnow() - dt_util.utc_from_timestamp(float(
+        return self.client.name or self.client.hostname
-                client.last_seen))) < self._detection_time
+
+    @property
+    def unique_id(self) -> str:
+        """Return a unique identifier for this client."""
+        return '{}-{}'.format(self.client.mac, self.controller.site)
+
+    @property
+    def available(self) -> bool:
+        """Return if controller is available."""
+        return self.controller.available
+
+    @property
+    def device_info(self):
+        """Return a device description for device registry."""
+        return {
+            'connections': {(CONNECTION_NETWORK_MAC, self.client.mac)}
         }
 
-    async def async_scan_devices(self):
+    @property
-        """Scan for devices."""
+    def device_state_attributes(self):
-        await self.async_update()
+        """Return the device state attributes."""
-        return self._clients.keys()
-
-    def get_device_name(self, device):
-        """Return the name (if known) of the device.
-
-        If a name has been set in Unifi, then return that, else
-        return the hostname if it has been detected.
-        """
-        client = self._clients.get(device, {})
-        name = client.get('name') or client.get('hostname')
-        _LOGGER.debug("Device mac %s name %s", device, name)
-        return name
-
-    def get_extra_attributes(self, device):
-        """Return the extra attributes of the device."""
-        if not self._monitored_conditions:
-            return {}
-
-        client = self._clients.get(device, {})
         attributes = {}
-        for variable in self._monitored_conditions:
-            if variable in client:
-                if variable in TIMESTAMP_ATTRS:
-                    attributes[variable] = dt_util.utc_from_timestamp(
-                        float(client[variable])
-                    )
-                else:
-                    attributes[variable] = client[variable]
 
-        _LOGGER.debug("Device mac %s attributes %s", device, attributes)
+        for variable in DEVICE_ATTRIBUTES:
+            if variable in self.client.raw:
+                attributes[variable] = self.client.raw[variable]
+
         return attributes

homeassistant/components/unifi/switch.py

@@ -1,5 +1,4 @@
 """Support for devices connected to UniFi POE."""
-from datetime import timedelta
 import logging
 
 from homeassistant.components import unifi
@@ -11,8 +10,6 @@ from homeassistant.helpers.dispatcher import async_dispatcher_connect
 
 from .const import CONF_CONTROLLER, CONF_SITE_ID, CONTROLLER_ID
 
-SCAN_INTERVAL = timedelta(seconds=15)
-
 LOGGER = logging.getLogger(__name__)
 
 
@@ -32,6 +29,10 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         site=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID],
     )
     controller = hass.data[unifi.DOMAIN][controller_id]
+
+    if controller.site_role != 'admin':
+        return
+
     switches = {}
 
     @callback

tests/components/unifi/test_controller.py

@@ -4,13 +4,22 @@ from unittest.mock import Mock, patch
 import pytest
 
 from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.components.unifi.const import CONF_CONTROLLER, CONF_SITE_ID
+from homeassistant.components.unifi.const import (
+    CONF_CONTROLLER, CONF_SITE_ID, UNIFI_CONFIG)
 from homeassistant.const import (
     CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME, CONF_VERIFY_SSL)
 from homeassistant.components.unifi import controller, errors
 
 from tests.common import mock_coro
 
+CONTROLLER_SITES = {
+    'site1': {
+        'desc': 'nice name',
+        'name': 'site',
+        'role': 'admin'
+    }
+}
+
 CONTROLLER_DATA = {
     CONF_HOST: '1.2.3.4',
     CONF_USERNAME: 'username',
@@ -28,10 +37,12 @@ ENTRY_CONFIG = {
 async def test_controller_setup():
     """Successful setup."""
     hass = Mock()
+    hass.data = {UNIFI_CONFIG: {}}
     entry = Mock()
     entry.data = ENTRY_CONFIG
     api = Mock()
     api.initialize.return_value = mock_coro(True)
+    api.sites.return_value = mock_coro(CONTROLLER_SITES)
 
     unifi_controller = controller.UniFiController(hass, entry)
 
@@ -40,8 +51,10 @@ async def test_controller_setup():
         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 len(hass.config_entries.async_forward_entry_setup.mock_calls) == 2
     assert hass.config_entries.async_forward_entry_setup.mock_calls[0][1] == \
+        (entry, 'device_tracker')
+    assert hass.config_entries.async_forward_entry_setup.mock_calls[1][1] == \
         (entry, 'switch')
 
 
@@ -53,12 +66,24 @@ async def test_controller_host():
 
     unifi_controller = controller.UniFiController(hass, entry)
 
-    assert unifi_controller.host == '1.2.3.4'
+    assert unifi_controller.host == CONTROLLER_DATA[CONF_HOST]
+
+
+async def test_controller_site():
+    """Config entry site and controller site are the same."""
+    hass = Mock()
+    entry = Mock()
+    entry.data = ENTRY_CONFIG
+
+    unifi_controller = controller.UniFiController(hass, entry)
+
+    assert unifi_controller.site == CONTROLLER_DATA[CONF_SITE_ID]
 
 
 async def test_controller_mac():
     """Test that it is possible to identify controller mac."""
     hass = Mock()
+    hass.data = {UNIFI_CONFIG: {}}
     entry = Mock()
     entry.data = ENTRY_CONFIG
     client = Mock()
@@ -67,6 +92,7 @@ async def test_controller_mac():
     api = Mock()
     api.initialize.return_value = mock_coro(True)
     api.clients = {'client1': client}
+    api.sites.return_value = mock_coro(CONTROLLER_SITES)
 
     unifi_controller = controller.UniFiController(hass, entry)
 
@@ -80,6 +106,7 @@ async def test_controller_mac():
 async def test_controller_no_mac():
     """Test that it works to not find the controllers mac."""
     hass = Mock()
+    hass.data = {UNIFI_CONFIG: {}}
     entry = Mock()
     entry.data = ENTRY_CONFIG
     client = Mock()
@@ -87,6 +114,7 @@ async def test_controller_no_mac():
     api = Mock()
     api.initialize.return_value = mock_coro(True)
     api.clients = {'client1': client}
+    api.sites.return_value = mock_coro(CONTROLLER_SITES)
 
     unifi_controller = controller.UniFiController(hass, entry)
 
@@ -149,10 +177,12 @@ async def test_reset_if_entry_had_wrong_auth():
 async def test_reset_unloads_entry_if_setup():
     """Calling reset when the entry has been setup."""
     hass = Mock()
+    hass.data = {UNIFI_CONFIG: {}}
     entry = Mock()
     entry.data = ENTRY_CONFIG
     api = Mock()
     api.initialize.return_value = mock_coro(True)
+    api.sites.return_value = mock_coro(CONTROLLER_SITES)
 
     unifi_controller = controller.UniFiController(hass, entry)
 
@@ -160,13 +190,13 @@ async def test_reset_unloads_entry_if_setup():
                       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
+    assert len(hass.config_entries.async_forward_entry_setup.mock_calls) == 2
 
     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
+    assert len(hass.config_entries.async_forward_entry_unload.mock_calls) == 2
 
 
 async def test_get_controller(hass):

tests/components/unifi/test_device_tracker.py

@@ -1,267 +1,160 @@
 """The tests for the Unifi WAP device tracker platform."""
-from unittest import mock
+from collections import deque
-from datetime import datetime, timedelta
+from copy import copy
+from unittest.mock import Mock
+from datetime import timedelta
 
 import pytest
-import voluptuous as vol
 
-import homeassistant.util.dt as dt_util
-from homeassistant.components.device_tracker import DOMAIN
-import homeassistant.components.unifi.device_tracker as unifi
-from homeassistant.const import (CONF_HOST, CONF_USERNAME, CONF_PASSWORD,
-                                 CONF_PLATFORM, CONF_VERIFY_SSL,
-                                 CONF_MONITORED_CONDITIONS)
-
-from tests.common import mock_coro
-from asynctest import CoroutineMock
 from aiounifi.clients import Clients
+from aiounifi.devices import Devices
+
+from homeassistant import config_entries
+from homeassistant.components import unifi
+from homeassistant.components.unifi.const import (
+    CONF_CONTROLLER, CONF_SITE_ID, UNIFI_CONFIG)
+from homeassistant.const import (
+    CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME, CONF_VERIFY_SSL)
+from homeassistant.setup import async_setup_component
+
+import homeassistant.components.device_tracker as device_tracker
+import homeassistant.components.unifi.device_tracker as unifi_dt
+import homeassistant.util.dt as dt_util
 
 DEFAULT_DETECTION_TIME = timedelta(seconds=300)
 
+CLIENT_1 = {
+    'essid': 'ssid',
+    'hostname': 'client_1',
+    'ip': '10.0.0.1',
+    'is_wired': False,
+    'last_seen': 1562600145,
+    'mac': '00:00:00:00:00:01',
+}
+CLIENT_2 = {
+    'hostname': 'client_2',
+    'ip': '10.0.0.2',
+    'is_wired': True,
+    'last_seen': 1562600145,
+    'mac': '00:00:00:00:00:02',
+    'name': 'Wired Client',
+}
+CLIENT_3 = {
+    'essid': 'ssid2',
+    'hostname': 'client_3',
+    'ip': '10.0.0.3',
+    'is_wired': False,
+    'last_seen': 1562600145,
+    'mac': '00:00:00:00:00:03',
+}
 
-@pytest.fixture
+CONTROLLER_DATA = {
-def mock_ctrl():
+    CONF_HOST: 'mock-host',
-    """Mock pyunifi."""
+    CONF_USERNAME: 'mock-user',
-    with mock.patch('aiounifi.Controller') as mock_control:
+    CONF_PASSWORD: 'mock-pswd',
-        mock_control.return_value.login.return_value = mock_coro()
+    CONF_PORT: 1234,
-        mock_control.return_value.initialize.return_value = mock_coro()
+    CONF_SITE_ID: 'mock-site',
-        yield mock_control
+    CONF_VERIFY_SSL: True
+}
+
+ENTRY_CONFIG = {
+    CONF_CONTROLLER: CONTROLLER_DATA
+}
+
+CONTROLLER_ID = unifi.CONTROLLER_ID.format(host='mock-host', site='mock-site')
 
 
 @pytest.fixture
-def mock_scanner():
+def mock_controller(hass):
-    """Mock UnifyScanner."""
+    """Mock a UniFi Controller."""
-    with mock.patch('homeassistant.components.unifi.device_tracker'
+    hass.data[UNIFI_CONFIG] = {}
-                    '.UnifiScanner') as scanner:
+    controller = unifi.UniFiController(hass, None)
-        yield scanner
+
+    controller.api = Mock()
+    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
 
 
-@mock.patch('os.access', return_value=True)
+async def setup_controller(hass, mock_controller):
-@mock.patch('os.path.isfile', mock.Mock(return_value=True))
+    """Load the UniFi switch platform with the provided controller."""
-async def test_config_valid_verify_ssl(hass, mock_scanner, mock_ctrl):
+    hass.config.components.add(unifi.DOMAIN)
-    """Test the setup with a string for ssl_verify.
+    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)
+    mock_controller.config_entry = config_entry
 
-    Representing the absolute path to a CA certificate bundle.
+    await mock_controller.async_update()
-    """
+    await hass.config_entries.async_forward_entry_setup(
-    config = {
+        config_entry, device_tracker.DOMAIN)
-        DOMAIN: unifi.PLATFORM_SCHEMA({
-            CONF_PLATFORM: unifi.DOMAIN,
-            CONF_USERNAME: 'foo',
-            CONF_PASSWORD: 'password',
-            CONF_VERIFY_SSL: "/tmp/unifi.crt"
-        })
-    }
-    result = await unifi.async_get_scanner(hass, config)
-    assert mock_scanner.return_value == result
-    assert mock_ctrl.call_count == 1
 
-    assert mock_scanner.call_count == 1
+    await hass.async_block_till_done()
-    assert mock_scanner.call_args == mock.call(mock_ctrl.return_value,
-                                               DEFAULT_DETECTION_TIME,
-                                               None, None)
 
 
-async def test_config_minimal(hass, mock_scanner, mock_ctrl):
+async def test_platform_manually_configured(hass):
-    """Test the setup with minimal configuration."""
+    """Test that we do not discover anything or try to set up a bridge."""
-    config = {
+    assert await async_setup_component(hass, device_tracker.DOMAIN, {
-        DOMAIN: unifi.PLATFORM_SCHEMA({
+        device_tracker.DOMAIN: {
-            CONF_PLATFORM: unifi.DOMAIN,
+            'platform': 'unifi'
-            CONF_USERNAME: 'foo',
-            CONF_PASSWORD: 'password',
-        })
-    }
-
-    result = await unifi.async_get_scanner(hass, config)
-    assert mock_scanner.return_value == result
-    assert mock_ctrl.call_count == 1
-
-    assert mock_scanner.call_count == 1
-    assert mock_scanner.call_args == mock.call(mock_ctrl.return_value,
-                                               DEFAULT_DETECTION_TIME,
-                                               None, None)
-
-
-async def test_config_full(hass, mock_scanner, mock_ctrl):
-    """Test the setup with full configuration."""
-    config = {
-        DOMAIN: unifi.PLATFORM_SCHEMA({
-            CONF_PLATFORM: unifi.DOMAIN,
-            CONF_USERNAME: 'foo',
-            CONF_PASSWORD: 'password',
-            CONF_HOST: 'myhost',
-            CONF_VERIFY_SSL: False,
-            CONF_MONITORED_CONDITIONS: ['essid', 'signal'],
-            'port': 123,
-            'site_id': 'abcdef01',
-            'detection_time': 300,
-        })
-    }
-    result = await unifi.async_get_scanner(hass, config)
-    assert mock_scanner.return_value == result
-    assert mock_ctrl.call_count == 1
-
-    assert mock_scanner.call_count == 1
-    assert mock_scanner.call_args == mock.call(
-        mock_ctrl.return_value,
-        DEFAULT_DETECTION_TIME,
-        None,
-        config[DOMAIN][CONF_MONITORED_CONDITIONS])
-
-
-def test_config_error():
-    """Test for configuration errors."""
-    with pytest.raises(vol.Invalid):
-        unifi.PLATFORM_SCHEMA({
-            # no username
-            CONF_PLATFORM: unifi.DOMAIN,
-            CONF_HOST: 'myhost',
-            'port': 123,
-        })
-    with pytest.raises(vol.Invalid):
-        unifi.PLATFORM_SCHEMA({
-            CONF_PLATFORM: unifi.DOMAIN,
-            CONF_USERNAME: 'foo',
-            CONF_PASSWORD: 'password',
-            CONF_HOST: 'myhost',
-            'port': 'foo',  # bad port!
-        })
-    with pytest.raises(vol.Invalid):
-        unifi.PLATFORM_SCHEMA({
-            CONF_PLATFORM: unifi.DOMAIN,
-            CONF_USERNAME: 'foo',
-            CONF_PASSWORD: 'password',
-            CONF_VERIFY_SSL: "dfdsfsdfsd",  # Invalid ssl_verify (no file)
-        })
-
-
-async def test_config_controller_failed(hass, mock_ctrl, mock_scanner):
-    """Test for controller failure."""
-    config = {
-        'device_tracker': {
-            CONF_PLATFORM: unifi.DOMAIN,
-            CONF_USERNAME: 'foo',
-            CONF_PASSWORD: 'password',
         }
-    }
+    }) is True
-    mock_ctrl.side_effect = unifi.CannotConnect
+    assert unifi.DOMAIN not in hass.data
-    result = await unifi.async_get_scanner(hass, config)
-    assert result is False
 
 
-async def test_scanner_update():
+async def test_no_clients(hass, mock_controller):
-    """Test the scanner update."""
+    """Test the update_clients function when no clients are found."""
-    ctrl = mock.MagicMock()
+    mock_controller.mock_client_responses.append({})
-    fake_clients = [
+    mock_controller.mock_device_responses.append({})
-        {'mac': '123', 'essid': 'barnet',
+    await setup_controller(hass, mock_controller)
-         'last_seen': dt_util.as_timestamp(dt_util.utcnow())},
+    assert len(mock_controller.mock_requests) == 2
-        {'mac': '234', 'essid': 'barnet',
+    assert len(hass.states.async_all()) == 2
-         'last_seen': dt_util.as_timestamp(dt_util.utcnow())},
-    ]
-    ctrl.clients = Clients([], CoroutineMock(return_value=fake_clients))
-    scnr = unifi.UnifiScanner(ctrl, DEFAULT_DETECTION_TIME, None, None)
-    await scnr.async_update()
-    assert len(scnr._clients) == 2
 
 
-def test_scanner_update_error():
+async def test_tracked_devices(hass, mock_controller):
-    """Test the scanner update for error."""
+    """Test the update_items function with some clients."""
-    ctrl = mock.MagicMock()
+    mock_controller.mock_client_responses.append(
-    ctrl.get_clients.side_effect = unifi.aiounifi.AiounifiException
+        [CLIENT_1, CLIENT_2, CLIENT_3])
-    unifi.UnifiScanner(ctrl, DEFAULT_DETECTION_TIME, None, None)
+    mock_controller.mock_device_responses.append({})
+    mock_controller.unifi_config = {unifi_dt.CONF_SSID_FILTER: ['ssid']}
 
+    await setup_controller(hass, mock_controller)
+    assert len(mock_controller.mock_requests) == 2
+    assert len(hass.states.async_all()) == 4
 
-async def test_scan_devices():
+    device_1 = hass.states.get('device_tracker.client_1')
-    """Test the scanning for devices."""
+    assert device_1 is not None
-    ctrl = mock.MagicMock()
+    assert device_1.state == 'not_home'
-    fake_clients = [
-        {'mac': '123', 'essid': 'barnet',
-         'last_seen': dt_util.as_timestamp(dt_util.utcnow())},
-        {'mac': '234', 'essid': 'barnet',
-         'last_seen': dt_util.as_timestamp(dt_util.utcnow())},
-    ]
-    ctrl.clients = Clients([], CoroutineMock(return_value=fake_clients))
-    scnr = unifi.UnifiScanner(ctrl, DEFAULT_DETECTION_TIME, None, None)
-    await scnr.async_update()
-    assert set(await scnr.async_scan_devices()) == set(['123', '234'])
 
+    device_2 = hass.states.get('device_tracker.wired_client')
+    assert device_2 is not None
+    assert device_2.state == 'not_home'
 
-async def test_scan_devices_filtered():
+    device_3 = hass.states.get('device_tracker.client_3')
-    """Test the scanning for devices based on SSID."""
+    assert device_3 is None
-    ctrl = mock.MagicMock()
-    fake_clients = [
-        {'mac': '123', 'essid': 'foonet',
-         'last_seen': dt_util.as_timestamp(dt_util.utcnow())},
-        {'mac': '234', 'essid': 'foonet',
-         'last_seen': dt_util.as_timestamp(dt_util.utcnow())},
-        {'mac': '567', 'essid': 'notnet',
-         'last_seen': dt_util.as_timestamp(dt_util.utcnow())},
-        {'mac': '890', 'essid': 'barnet',
-         'last_seen': dt_util.as_timestamp(dt_util.utcnow())},
-    ]
 
-    ssid_filter = ['foonet', 'barnet']
+    client_1 = copy(CLIENT_1)
-    ctrl.clients = Clients([], CoroutineMock(return_value=fake_clients))
+    client_1['last_seen'] = dt_util.as_timestamp(dt_util.utcnow())
-    scnr = unifi.UnifiScanner(ctrl, DEFAULT_DETECTION_TIME, ssid_filter, None)
+    mock_controller.mock_client_responses.append([client_1])
-    await scnr.async_update()
+    mock_controller.mock_device_responses.append({})
-    assert set(await scnr.async_scan_devices()) == set(['123', '234', '890'])
+    await mock_controller.async_update()
+    await hass.async_block_till_done()
 
-
+    device_1 = hass.states.get('device_tracker.client_1')
-async def test_get_device_name():
+    assert device_1.state == 'home'
-    """Test the getting of device names."""
-    ctrl = mock.MagicMock()
-    fake_clients = [
-        {'mac': '123',
-         'hostname': 'foobar',
-         'essid': 'barnet',
-         'last_seen': dt_util.as_timestamp(dt_util.utcnow())},
-        {'mac': '234',
-         'name': 'Nice Name',
-         'essid': 'barnet',
-         'last_seen': dt_util.as_timestamp(dt_util.utcnow())},
-        {'mac': '456',
-         'essid': 'barnet',
-         'last_seen': '1504786810'},
-    ]
-    ctrl.clients = Clients([], CoroutineMock(return_value=fake_clients))
-    scnr = unifi.UnifiScanner(ctrl, DEFAULT_DETECTION_TIME, None, None)
-    await scnr.async_update()
-    assert scnr.get_device_name('123') == 'foobar'
-    assert scnr.get_device_name('234') == 'Nice Name'
-    assert scnr.get_device_name('456') is None
-    assert scnr.get_device_name('unknown') is None
-
-
-async def test_monitored_conditions():
-    """Test the filtering of attributes."""
-    ctrl = mock.MagicMock()
-    fake_clients = [
-        {'mac': '123',
-         'hostname': 'foobar',
-         'essid': 'barnet',
-         'signal': -60,
-         'last_seen': dt_util.as_timestamp(dt_util.utcnow()),
-         'latest_assoc_time': 946684800.0},
-        {'mac': '234',
-         'name': 'Nice Name',
-         'essid': 'barnet',
-         'signal': -42,
-         'last_seen': dt_util.as_timestamp(dt_util.utcnow())},
-        {'mac': '456',
-         'hostname': 'wired',
-         'essid': 'barnet',
-         'last_seen': dt_util.as_timestamp(dt_util.utcnow())},
-    ]
-    ctrl.clients = Clients([], CoroutineMock(return_value=fake_clients))
-    scnr = unifi.UnifiScanner(ctrl, DEFAULT_DETECTION_TIME, None,
-                              ['essid', 'signal', 'latest_assoc_time'])
-    await scnr.async_update()
-    assert scnr.get_extra_attributes('123') == {
-        'essid': 'barnet',
-        'signal': -60,
-        'latest_assoc_time': datetime(2000, 1, 1, 0, 0, tzinfo=dt_util.UTC)
-    }
-    assert scnr.get_extra_attributes('234') == {
-        'essid': 'barnet',
-        'signal': -42
-    }
-    assert scnr.get_extra_attributes('456') == {'essid': 'barnet'}

tests/components/unifi/test_init.py

@@ -1,4 +1,5 @@
 """Test UniFi setup process."""
+from datetime import timedelta
 from unittest.mock import Mock, patch
 
 from homeassistant.components import unifi
@@ -15,6 +16,29 @@ 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
+    assert hass.data[unifi.UNIFI_CONFIG] == []
+
+
+async def test_setup_with_config(hass):
+    """Test that we do not discover anything or try to set up a bridge."""
+    config = {
+        unifi.DOMAIN: {
+            unifi.CONF_CONTROLLERS: {
+                unifi.CONF_HOST: '1.2.3.4',
+                unifi.CONF_SITE_ID: 'My site',
+                unifi.CONF_DETECTION_TIME: 3,
+                unifi.CONF_SSID_FILTER: ['ssid']
+            }
+        }
+    }
+    assert await async_setup_component(hass, unifi.DOMAIN, config) is True
+    assert unifi.DOMAIN not in hass.data
+    assert hass.data[unifi.UNIFI_CONFIG] == [{
+        unifi.CONF_HOST: '1.2.3.4',
+        unifi.CONF_SITE_ID: 'My site',
+        unifi.CONF_DETECTION_TIME: timedelta(seconds=3),
+        unifi.CONF_SSID_FILTER: ['ssid']
+    }]
 
 
 async def test_successful_config_entry(hass):
@@ -247,41 +271,6 @@ async def test_controller_site_already_configured(hass):
     assert result['type'] == 'abort'
 
 
-async def test_user_permissions_low(hass, aioclient_mock):
-    """Test config flow."""
-    flow = config_flow.UnifiFlowHandler()
-    flow.hass = hass
-
-    with patch('aiounifi.Controller') as mock_controller:
-        def mock_constructor(
-                host, username, password, port, site, websession, sslcontext):
-            """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={
-            CONF_HOST: '1.2.3.4',
-            CONF_USERNAME: 'username',
-            CONF_PASSWORD: 'password',
-            CONF_PORT: 1234,
-            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 = config_flow.UnifiFlowHandler()

tests/components/unifi/test_switch.py

@@ -12,7 +12,8 @@ from aiounifi.devices import Devices
 
 from homeassistant import config_entries
 from homeassistant.components import unifi
-from homeassistant.components.unifi.const import CONF_CONTROLLER, CONF_SITE_ID
+from homeassistant.components.unifi.const import (
+    CONF_CONTROLLER, CONF_SITE_ID, UNIFI_CONFIG)
 from homeassistant.setup import async_setup_component
 from homeassistant.const import (
     CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME, CONF_VERIFY_SSL)
@@ -188,8 +189,11 @@ CONTROLLER_ID = unifi.CONTROLLER_ID.format(host='mock-host', site='mock-site')
 @pytest.fixture
 def mock_controller(hass):
     """Mock a UniFi Controller."""
+    hass.data[UNIFI_CONFIG] = {}
     controller = unifi.UniFiController(hass, None)
 
+    controller._site_role = 'admin'
+
     controller.api = Mock()
     controller.mock_requests = []
 
@@ -257,13 +261,25 @@ async def test_controller_not_client(hass, mock_controller):
     assert cloudkey is None
 
 
-async def test_switches(hass, mock_controller):
+async def test_not_admin(hass, mock_controller):
-    """Test the update_items function with some lights."""
+    """Test that switch platform only work on an admin account."""
-    mock_controller.mock_client_responses.append([CLIENT_1, CLIENT_4])
+    mock_controller.mock_client_responses.append([CLIENT_1])
-    mock_controller.mock_device_responses.append([DEVICE_1])
+    mock_controller.mock_device_responses.append([])
+
+    mock_controller._site_role = 'viewer'
+
+    await setup_controller(hass, mock_controller)
+    assert len(mock_controller.mock_requests) == 2
+    assert len(hass.states.async_all()) == 0
+
+
+async def test_switches(hass, mock_controller):
+    """Test the update_items function with some clients."""
+    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')
@@ -276,8 +292,8 @@ async def test_switches(hass, mock_controller):
     assert switch_1.attributes['port'] == 1
     assert switch_1.attributes['poe_mode'] == 'auto'
 
-    switch = hass.states.get('switch.client_4')
+    switch_4 = hass.states.get('switch.client_4')
-    assert switch is None
+    assert switch_4 is None
 
 
 async def test_new_client_discovered(hass, mock_controller):
@@ -296,13 +312,37 @@ async def test_new_client_discovered(hass, mock_controller):
     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
+    assert mock_controller.mock_requests[2] == {
+        'json': {
+            'port_overrides': [{
+                'port_idx': 1,
+                'portconf_id': '1a1',
+                'poe_mode': 'off'}]
+        },
+        'method': 'put',
+        'path': 's/{site}/rest/device/mock-id'
+    }
 
-    switch = hass.states.get('switch.client_2')
+    await hass.services.async_call('switch', 'turn_on', {
-    assert switch is not None
+        'entity_id': 'switch.client_1'
-    assert switch.state == 'on'
+    }, blocking=True)
+    assert len(mock_controller.mock_requests) == 7
+    assert mock_controller.mock_requests[5] == {
+        'json': {
+            'port_overrides': [{
+                'port_idx': 1,
+                'portconf_id': '1a1',
+                'poe_mode': 'auto'}]
+        },
+        'method': 'put',
+        'path': 's/{site}/rest/device/mock-id'
+    }
+
+    switch_2 = hass.states.get('switch.client_2')
+    assert switch_2 is not None
+    assert switch_2.state == 'on'
 
 
 async def test_failed_update_successful_login(hass, mock_controller):
@@ -346,7 +386,7 @@ async def test_failed_update_unreachable_controller(hass, mock_controller):
     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