UniFi - Improve controller tests (#27261)

* Improve controller tests and harmonize setup_unifi_integration to one
* Store listeners to dispatchers to be used during reset

homeassistant/components/unifi/__init__.py

@@ -77,12 +77,13 @@ async def async_setup_entry(hass, config_entry):
         hass.data[DOMAIN] = {}
 
     controller = UniFiController(hass, config_entry)
-    controller_id = get_controller_id_from_config_entry(config_entry)
-    hass.data[DOMAIN][controller_id] = controller
 
     if not await controller.async_setup():
         return False
 
+    controller_id = get_controller_id_from_config_entry(config_entry)
+    hass.data[DOMAIN][controller_id] = controller
+
     if controller.mac is None:
         return True
 

homeassistant/components/unifi/controller.py

@@ -57,6 +57,7 @@ class UniFiController:
         self.progress = None
         self.wireless_clients = None
 
+        self.listeners = []
         self._site_name = None
         self._site_role = None
 
@@ -258,13 +259,14 @@ class UniFiController:
 
     def import_configuration(self):
         """Import configuration to config entry options."""
-        unifi_config = {}
+        import_config = {}
+
         for config in self.hass.data[UNIFI_CONFIG]:
             if (
                 self.host == config[CONF_HOST]
                 and self.site_name == config[CONF_SITE_ID]
             ):
-                unifi_config = config
+                import_config = config
                 break
 
         old_options = dict(self.config_entry.options)
@@ -278,16 +280,17 @@ class UniFiController:
             (CONF_DETECTION_TIME, CONF_DETECTION_TIME),
             (CONF_SSID_FILTER, CONF_SSID_FILTER),
         ):
-            if config in unifi_config:
+            if config in import_config:
-                if config == option and unifi_config[
+                print(config)
+                if config == option and import_config[
                     config
                 ] != self.config_entry.options.get(option):
-                    new_options[option] = unifi_config[config]
+                    new_options[option] = import_config[config]
                 elif config != option and (
                     option not in self.config_entry.options
-                    or unifi_config[config] == self.config_entry.options.get(option)
+                    or import_config[config] == self.config_entry.options.get(option)
                 ):
-                    new_options[option] = not unifi_config[config]
+                    new_options[option] = not import_config[config]
 
         if new_options:
             options = {**old_options, **new_options}
@@ -301,15 +304,15 @@ class UniFiController:
         Will cancel any scheduled setup retry and will unload
         the config entry.
         """
-        # If the authentication was wrong.
-        if self.api is None:
-            return True
-
         for platform in SUPPORTED_PLATFORMS:
             await self.hass.config_entries.async_forward_entry_unload(
                 self.config_entry, platform
             )
 
+        for unsub_dispatcher in self.listeners:
+            unsub_dispatcher()
+        self.listeners = []
+
         return True
 
 

homeassistant/components/unifi/device_tracker.py

@@ -67,7 +67,9 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         """Update the values of the controller."""
         update_items(controller, async_add_entities, tracked)
 
-    async_dispatcher_connect(hass, controller.signal_update, update_controller)
+    controller.listeners.append(
+        async_dispatcher_connect(hass, controller.signal_update, update_controller)
+    )
 
     @callback
     def update_disable_on_entities():
@@ -82,8 +84,10 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
                 entity.registry_entry.entity_id, disabled_by=disabled_by
             )
 
-    async_dispatcher_connect(
+    controller.listeners.append(
-        hass, controller.signal_options_update, update_disable_on_entities
+        async_dispatcher_connect(
+            hass, controller.signal_options_update, update_disable_on_entities
+        )
     )
 
     update_controller()

homeassistant/components/unifi/sensor.py

@@ -31,7 +31,9 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         """Update the values of the controller."""
         update_items(controller, async_add_entities, sensors)
 
-    async_dispatcher_connect(hass, controller.signal_update, update_controller)
+    controller.listeners.append(
+        async_dispatcher_connect(hass, controller.signal_update, update_controller)
+    )
 
     @callback
     def update_disable_on_entities():
@@ -46,8 +48,10 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
                 entity.registry_entry.entity_id, disabled_by=disabled_by
             )
 
-    async_dispatcher_connect(
+    controller.listeners.append(
-        hass, controller.signal_options_update, update_disable_on_entities
+        async_dispatcher_connect(
+            hass, controller.signal_options_update, update_disable_on_entities
+        )
     )
 
     update_controller()

homeassistant/components/unifi/switch.py

@@ -53,7 +53,9 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         """Update the values of the controller."""
         update_items(controller, async_add_entities, switches, switches_off)
 
-    async_dispatcher_connect(hass, controller.signal_update, update_controller)
+    controller.listeners.append(
+        async_dispatcher_connect(hass, controller.signal_update, update_controller)
+    )
 
     update_controller()
     switches_off.clear()

tests/components/unifi/test_controller.py

@@ -1,9 +1,14 @@
 """Test UniFi Controller."""
-from unittest.mock import Mock, patch
+from collections import deque
+from datetime import timedelta
+
+from asynctest import Mock, patch
 
 import pytest
 
+from homeassistant import config_entries
 from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.components import unifi
 from homeassistant.components.unifi.const import (
     CONF_CONTROLLER,
     CONF_SITE_ID,
@@ -17,269 +22,362 @@ from homeassistant.const import (
     CONF_USERNAME,
     CONF_VERIFY_SSL,
 )
-from homeassistant.components.unifi import controller, errors
+import aiounifi
 
-from tests.common import mock_coro
+CONTROLLER_HOST = {
-
+    "hostname": "controller_host",
-CONTROLLER_SITES = {"site1": {"desc": "nice name", "name": "site", "role": "admin"}}
+    "ip": "1.2.3.4",
+    "is_wired": True,
+    "last_seen": 1562600145,
+    "mac": "10:00:00:00:00:01",
+    "name": "Controller host",
+    "oui": "Producer",
+    "sw_mac": "00:00:00:00:01:01",
+    "sw_port": 1,
+    "wired-rx_bytes": 1234000000,
+    "wired-tx_bytes": 5678000000,
+}
 
 CONTROLLER_DATA = {
     CONF_HOST: "1.2.3.4",
     CONF_USERNAME: "username",
     CONF_PASSWORD: "password",
     CONF_PORT: 1234,
-    CONF_SITE_ID: "site",
+    CONF_SITE_ID: "site_id",
-    CONF_VERIFY_SSL: True,
+    CONF_VERIFY_SSL: False,
 }
 
 ENTRY_CONFIG = {CONF_CONTROLLER: CONTROLLER_DATA}
 
+SITES = {"Site name": {"desc": "Site name", "name": "site_id", "role": "admin"}}
 
-async def test_controller_setup():
+
+async def setup_unifi_integration(
+    hass,
+    config,
+    options,
+    sites,
+    clients_response,
+    devices_response,
+    clients_all_response,
+):
+    """Create the UniFi controller."""
+    if UNIFI_CONFIG not in hass.data:
+        hass.data[UNIFI_CONFIG] = []
+    hass.data[UNIFI_WIRELESS_CLIENTS] = unifi.UnifiWirelessClients(hass)
+    config_entry = config_entries.ConfigEntry(
+        version=1,
+        domain=unifi.DOMAIN,
+        title="Mock Title",
+        data=config,
+        source="test",
+        connection_class=config_entries.CONN_CLASS_LOCAL_POLL,
+        system_options={},
+        options=options,
+        entry_id=1,
+    )
+
+    mock_client_responses = deque()
+    mock_client_responses.append(clients_response)
+
+    mock_device_responses = deque()
+    mock_device_responses.append(devices_response)
+
+    mock_client_all_responses = deque()
+    mock_client_all_responses.append(clients_all_response)
+
+    mock_requests = []
+
+    async def mock_request(self, method, path, json=None):
+        mock_requests.append({"method": method, "path": path, "json": json})
+
+        if path == "s/{site}/stat/sta" and mock_client_responses:
+            return mock_client_responses.popleft()
+        if path == "s/{site}/stat/device" and mock_device_responses:
+            return mock_device_responses.popleft()
+        if path == "s/{site}/rest/user" and mock_client_all_responses:
+            return mock_client_all_responses.popleft()
+        return {}
+
+    with patch("aiounifi.Controller.login", return_value=True), patch(
+        "aiounifi.Controller.sites", return_value=sites
+    ), patch("aiounifi.Controller.request", new=mock_request):
+        await unifi.async_setup_entry(hass, config_entry)
+    await hass.async_block_till_done()
+    hass.config_entries._entries.append(config_entry)
+
+    controller_id = unifi.get_controller_id_from_config_entry(config_entry)
+    if controller_id not in hass.data[unifi.DOMAIN]:
+        return None
+    controller = hass.data[unifi.DOMAIN][controller_id]
+
+    controller.mock_client_responses = mock_client_responses
+    controller.mock_device_responses = mock_device_responses
+    controller.mock_client_all_responses = mock_client_all_responses
+    controller.mock_requests = mock_requests
+
+    return controller
+
+
+async def test_controller_setup(hass):
     """Successful setup."""
-    hass = Mock()
+    with patch(
-    hass.data = {
+        "homeassistant.config_entries.ConfigEntries.async_forward_entry_setup",
-        UNIFI_CONFIG: [
+        return_value=True,
-            {
+    ) as forward_entry_setup:
-                CONF_HOST: CONTROLLER_DATA[CONF_HOST],
+        controller = await setup_unifi_integration(
-                CONF_SITE_ID: "nice name",
+            hass,
-                controller.CONF_BLOCK_CLIENT: ["mac"],
+            ENTRY_CONFIG,
-                controller.CONF_DONT_TRACK_CLIENTS: True,
+            options={},
-                controller.CONF_DONT_TRACK_DEVICES: True,
+            sites=SITES,
-                controller.CONF_DONT_TRACK_WIRED_CLIENTS: True,
+            clients_response=[],
-                controller.CONF_DETECTION_TIME: 30,
+            devices_response=[],
-                controller.CONF_SSID_FILTER: ["ssid"],
+            clients_all_response=[],
-            }
+        )
-        ],
-        UNIFI_WIRELESS_CLIENTS: Mock(),
-    }
-    entry = Mock()
-    entry.data = ENTRY_CONFIG
-    entry.options = {}
-    api = Mock()
-    api.initialize.return_value = mock_coro(True)
-    api.sites.return_value = mock_coro(CONTROLLER_SITES)
-    api.clients = []
 
-    unifi_controller = controller.UniFiController(hass, entry)
+        entry = controller.config_entry
+        assert len(forward_entry_setup.mock_calls) == len(
+            unifi.controller.SUPPORTED_PLATFORMS
+        )
+        assert forward_entry_setup.mock_calls[0][1] == (entry, "device_tracker")
+        assert forward_entry_setup.mock_calls[1][1] == (entry, "sensor")
+        assert forward_entry_setup.mock_calls[2][1] == (entry, "switch")
 
-    with patch.object(controller, "get_controller", return_value=mock_coro(api)):
+    assert controller.host == CONTROLLER_DATA[CONF_HOST]
-        assert await unifi_controller.async_setup() is True
+    assert controller.site == CONTROLLER_DATA[CONF_SITE_ID]
+    assert controller.site_name in SITES
+    assert controller.site_role == SITES[controller.site_name]["role"]
 
-    assert unifi_controller.api is api
+    assert (
-    assert len(hass.config_entries.async_forward_entry_setup.mock_calls) == len(
+        controller.option_allow_bandwidth_sensors
-        controller.SUPPORTED_PLATFORMS
+        == unifi.const.DEFAULT_ALLOW_BANDWIDTH_SENSORS
     )
-    assert hass.config_entries.async_forward_entry_setup.mock_calls[0][1] == (
+    assert controller.option_block_clients == unifi.const.DEFAULT_BLOCK_CLIENTS
-        entry,
+    assert controller.option_track_clients == unifi.const.DEFAULT_TRACK_CLIENTS
-        "device_tracker",
+    assert controller.option_track_devices == unifi.const.DEFAULT_TRACK_DEVICES
+    assert (
+        controller.option_track_wired_clients == unifi.const.DEFAULT_TRACK_WIRED_CLIENTS
     )
-    assert hass.config_entries.async_forward_entry_setup.mock_calls[1][1] == (
+    assert controller.option_detection_time == timedelta(
-        entry,
+        seconds=unifi.const.DEFAULT_DETECTION_TIME
-        "sensor",
-    )
-    assert hass.config_entries.async_forward_entry_setup.mock_calls[2][1] == (
-        entry,
-        "switch",
     )
+    assert controller.option_ssid_filter == unifi.const.DEFAULT_SSID_FILTER
+
+    assert controller.mac is None
+
+    assert controller.signal_update == "unifi-update-1.2.3.4-site_id"
+    assert controller.signal_options_update == "unifi-options-1.2.3.4-site_id"
 
 
-async def test_controller_host():
+async def test_controller_mac(hass):
-    """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 == 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()
+    controller = await setup_unifi_integration(
-    hass.data = {UNIFI_CONFIG: {}, UNIFI_WIRELESS_CLIENTS: Mock()}
+        hass,
-    hass.data[UNIFI_WIRELESS_CLIENTS].get_data.return_value = set()
+        ENTRY_CONFIG,
-    entry = Mock()
+        options={},
-    entry.data = ENTRY_CONFIG
+        sites=SITES,
-    entry.options = {}
+        clients_response=[CONTROLLER_HOST],
-    client = Mock()
+        devices_response=[],
-    client.ip = "1.2.3.4"
+        clients_all_response=[],
-    client.mac = "00:11:22:33:44:55"
+    )
-    api = Mock()
+    assert controller.mac == "10:00:00:00:00:01"
-    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)
-
-    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():
+async def test_controller_import_config(hass):
-    """Test that it works to not find the controllers mac."""
+    """Test that import configuration.yaml instructions work."""
-    hass = Mock()
+    hass.data[UNIFI_CONFIG] = [
-    hass.data = {UNIFI_CONFIG: {}, UNIFI_WIRELESS_CLIENTS: Mock()}
+        {
-    entry = Mock()
+            CONF_HOST: "1.2.3.4",
-    entry.data = ENTRY_CONFIG
+            CONF_SITE_ID: "Site name",
-    entry.options = {}
+            unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True,
-    client = Mock()
+            unifi.CONF_BLOCK_CLIENT: ["random mac"],
-    client.ip = "5.6.7.8"
+            unifi.CONF_DONT_TRACK_CLIENTS: True,
-    api = Mock()
+            unifi.CONF_DONT_TRACK_DEVICES: True,
-    api.initialize.return_value = mock_coro(True)
+            unifi.CONF_DONT_TRACK_WIRED_CLIENTS: True,
-    api.clients = {"client1": client}
+            unifi.CONF_DETECTION_TIME: 150,
-    api.sites.return_value = mock_coro(CONTROLLER_SITES)
+            unifi.CONF_SSID_FILTER: ["SSID"],
-    api.clients = {}
+        }
+    ]
+    controller = await setup_unifi_integration(
+        hass,
+        ENTRY_CONFIG,
+        options={},
+        sites=SITES,
+        clients_response=[],
+        devices_response=[],
+        clients_all_response=[],
+    )
 
-    unifi_controller = controller.UniFiController(hass, entry)
+    assert controller.option_allow_bandwidth_sensors is False
-
+    assert controller.option_block_clients == ["random mac"]
-    with patch.object(controller, "get_controller", return_value=mock_coro(api)):
+    assert controller.option_track_clients is False
-        assert await unifi_controller.async_setup() is True
+    assert controller.option_track_devices is False
-
+    assert controller.option_track_wired_clients is False
-    assert unifi_controller.mac is None
+    assert controller.option_detection_time == timedelta(seconds=150)
+    assert controller.option_ssid_filter == ["SSID"]
 
 
-async def test_controller_not_accessible():
+async def test_controller_not_accessible(hass):
     """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
+        unifi.controller, "get_controller", side_effect=unifi.errors.CannotConnect
     ), pytest.raises(ConfigEntryNotReady):
-        await unifi_controller.async_setup()
+        await setup_unifi_integration(
+            hass,
+            ENTRY_CONFIG,
+            options={},
+            sites=SITES,
+            clients_response=[],
+            devices_response=[],
+            clients_all_response=[],
+        )
 
 
-async def test_controller_unknown_error():
+async def test_controller_unknown_error(hass):
     """Unknown errors are handled."""
-    hass = Mock()
+    with patch.object(unifi.controller, "get_controller", side_effect=Exception):
-    entry = Mock()
+        await setup_unifi_integration(
-    entry.data = ENTRY_CONFIG
+            hass,
-    api = Mock()
+            ENTRY_CONFIG,
-    api.initialize.return_value = mock_coro(True)
+            options={},
-
+            sites=SITES,
-    unifi_controller = controller.UniFiController(hass, entry)
+            clients_response=[],
-
+            devices_response=[],
-    with patch.object(controller, "get_controller", side_effect=Exception):
+            clients_all_response=[],
-        assert await unifi_controller.async_setup() is False
+        )
-
+        assert hass.data[unifi.DOMAIN] == {}
-    assert not hass.helpers.event.async_call_later.mock_calls
 
 
-async def test_reset_if_entry_had_wrong_auth():
+async def test_reset_after_successful_setup(hass):
-    """Calling reset when the entry contains wrong auth."""
+    """Calling reset when the entry has been setup."""
-    hass = Mock()
+    controller = await setup_unifi_integration(
-    entry = Mock()
+        hass,
-    entry.data = ENTRY_CONFIG
+        ENTRY_CONFIG,
+        options={},
+        sites=SITES,
+        clients_response=[],
+        devices_response=[],
+        clients_all_response=[],
+    )
 
-    unifi_controller = controller.UniFiController(hass, entry)
+    assert len(controller.listeners) == 5
+
+    result = await controller.async_reset()
+    await hass.async_block_till_done()
+
+    assert result is True
+    assert len(controller.listeners) == 0
+
+
+async def test_failed_update_failed_login(hass):
+    """Running update can handle a failed login."""
+    controller = await setup_unifi_integration(
+        hass,
+        ENTRY_CONFIG,
+        options={},
+        sites=SITES,
+        clients_response=[],
+        devices_response=[],
+        clients_all_response=[],
+    )
 
     with patch.object(
-        controller, "get_controller", side_effect=errors.AuthenticationRequired
+        controller.api.clients, "update", side_effect=aiounifi.LoginRequired
+    ), patch.object(controller.api, "login", side_effect=aiounifi.AiounifiException):
+        await controller.async_update()
+    await hass.async_block_till_done()
+
+    assert controller.available is False
+
+
+async def test_failed_update_successful_login(hass):
+    """Running update can login when requested."""
+    controller = await setup_unifi_integration(
+        hass,
+        ENTRY_CONFIG,
+        options={},
+        sites=SITES,
+        clients_response=[],
+        devices_response=[],
+        clients_all_response=[],
+    )
+
+    with patch.object(
+        controller.api.clients, "update", side_effect=aiounifi.LoginRequired
+    ), patch.object(controller.api, "login", return_value=Mock(True)):
+        await controller.async_update()
+    await hass.async_block_till_done()
+
+    assert controller.available is True
+
+
+async def test_failed_update(hass):
+    """Running update can login when requested."""
+    controller = await setup_unifi_integration(
+        hass,
+        ENTRY_CONFIG,
+        options={},
+        sites=SITES,
+        clients_response=[],
+        devices_response=[],
+        clients_all_response=[],
+    )
+
+    with patch.object(
+        controller.api.clients, "update", side_effect=aiounifi.AiounifiException
     ):
-        assert await unifi_controller.async_setup() is False
+        await controller.async_update()
+    await hass.async_block_till_done()
 
-    assert not hass.async_add_job.mock_calls
+    assert controller.available is False
 
-    assert await unifi_controller.async_reset()
+    await controller.async_update()
-
+    await hass.async_block_till_done()
-
+    assert controller.available is True
-async def test_reset_unloads_entry_if_setup():
-    """Calling reset when the entry has been setup."""
-    hass = Mock()
-    hass.data = {UNIFI_CONFIG: {}, UNIFI_WIRELESS_CLIENTS: Mock()}
-    entry = Mock()
-    entry.data = ENTRY_CONFIG
-    entry.options = {}
-    api = Mock()
-    api.initialize.return_value = mock_coro(True)
-    api.sites.return_value = mock_coro(CONTROLLER_SITES)
-    api.clients = []
-
-    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) == len(
-        controller.SUPPORTED_PLATFORMS
-    )
-
-    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) == len(
-        controller.SUPPORTED_PLATFORMS
-    )
 
 
 async def test_get_controller(hass):
     """Successful call."""
-    with patch("aiounifi.Controller.login", return_value=mock_coro()):
+    with patch("aiounifi.Controller.login", return_value=Mock()):
-        assert await controller.get_controller(hass, **CONTROLLER_DATA)
+        assert await unifi.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[CONF_VERIFY_SSL] = False
-    with patch("aiounifi.Controller.login", return_value=mock_coro()):
+    with patch("aiounifi.Controller.login", return_value=Mock()):
-        assert await controller.get_controller(hass, **controller_data)
+        assert await unifi.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)
+            result = await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
-        except errors.AuthenticationRequired:
+        except unifi.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)
+            result = await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
-        except errors.CannotConnect:
+        except unifi.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)
+            result = await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
-        except errors.AuthenticationRequired:
+        except unifi.errors.AuthenticationRequired:
             pass
         assert result is None

tests/components/unifi/test_device_tracker.py

@@ -1,37 +1,23 @@
 """The tests for the UniFi device tracker platform."""
-from collections import deque
 from copy import copy
-
 from datetime import timedelta
 
-from asynctest import patch
-
 from homeassistant import config_entries
 from homeassistant.components import unifi
 from homeassistant.components.unifi.const import (
-    CONF_CONTROLLER,
-    CONF_SITE_ID,
     CONF_SSID_FILTER,
     CONF_TRACK_DEVICES,
     CONF_TRACK_WIRED_CLIENTS,
-    CONTROLLER_ID as CONF_CONTROLLER_ID,
-    UNIFI_CONFIG,
-    UNIFI_WIRELESS_CLIENTS,
-)
-from homeassistant.const import (
-    CONF_HOST,
-    CONF_PASSWORD,
-    CONF_PORT,
-    CONF_USERNAME,
-    CONF_VERIFY_SSL,
-    STATE_UNAVAILABLE,
 )
+from homeassistant.const import STATE_UNAVAILABLE
 from homeassistant.helpers import entity_registry
 from homeassistant.setup import async_setup_component
 
 import homeassistant.components.device_tracker as device_tracker
 import homeassistant.util.dt as dt_util
 
+from .test_controller import ENTRY_CONFIG, SITES, setup_unifi_integration
+
 DEFAULT_DETECTION_TIME = timedelta(seconds=300)
 
 CLIENT_1 = {
@@ -88,77 +74,6 @@ DEVICE_2 = {
     "version": "4.0.42.10433",
 }
 
-CONTROLLER_DATA = {
-    CONF_HOST: "mock-host",
-    CONF_USERNAME: "mock-user",
-    CONF_PASSWORD: "mock-pswd",
-    CONF_PORT: 1234,
-    CONF_SITE_ID: "mock-site",
-    CONF_VERIFY_SSL: False,
-}
-
-ENTRY_CONFIG = {CONF_CONTROLLER: CONTROLLER_DATA}
-
-CONTROLLER_ID = CONF_CONTROLLER_ID.format(host="mock-host", site="mock-site")
-
-
-async def setup_unifi_integration(
-    hass, config, options, clients_response, devices_response, clients_all_response
-):
-    """Create the UniFi controller."""
-    hass.data[UNIFI_CONFIG] = []
-    hass.data[UNIFI_WIRELESS_CLIENTS] = unifi.UnifiWirelessClients(hass)
-    config_entry = config_entries.ConfigEntry(
-        version=1,
-        domain=unifi.DOMAIN,
-        title="Mock Title",
-        data=config,
-        source="test",
-        connection_class=config_entries.CONN_CLASS_LOCAL_POLL,
-        system_options={},
-        options=options,
-        entry_id=1,
-    )
-
-    sites = {"Site name": {"desc": "Site name", "name": "mock-site", "role": "viewer"}}
-
-    mock_client_responses = deque()
-    mock_client_responses.append(clients_response)
-
-    mock_device_responses = deque()
-    mock_device_responses.append(devices_response)
-
-    mock_client_all_responses = deque()
-    mock_client_all_responses.append(clients_all_response)
-
-    mock_requests = []
-
-    async def mock_request(self, method, path, json=None):
-        mock_requests.append({"method": method, "path": path, "json": json})
-        if path == "s/{site}/stat/sta" and mock_client_responses:
-            return mock_client_responses.popleft()
-        if path == "s/{site}/stat/device" and mock_device_responses:
-            return mock_device_responses.popleft()
-        if path == "s/{site}/rest/user" and mock_client_all_responses:
-            return mock_client_all_responses.popleft()
-        return {}
-
-    with patch("aiounifi.Controller.login", return_value=True), patch(
-        "aiounifi.Controller.sites", return_value=sites
-    ), patch("aiounifi.Controller.request", new=mock_request):
-        await unifi.async_setup_entry(hass, config_entry)
-    await hass.async_block_till_done()
-    hass.config_entries._entries.append(config_entry)
-
-    controller_id = unifi.get_controller_id_from_config_entry(config_entry)
-    controller = hass.data[unifi.DOMAIN][controller_id]
-
-    controller.mock_client_responses = mock_client_responses
-    controller.mock_device_responses = mock_device_responses
-    controller.mock_client_all_responses = mock_client_all_responses
-
-    return controller
-
 
 async def test_platform_manually_configured(hass):
     """Test that nothing happens when configuring unifi through device tracker platform."""
@@ -177,9 +92,10 @@ async def test_no_clients(hass):
         hass,
         ENTRY_CONFIG,
         options={},
-        clients_response={},
+        sites=SITES,
-        devices_response={},
+        clients_response=[],
-        clients_all_response={},
+        devices_response=[],
+        clients_all_response=[],
     )
 
     assert len(hass.states.async_all()) == 2
@@ -191,6 +107,7 @@ async def test_tracked_devices(hass):
         hass,
         ENTRY_CONFIG,
         options={CONF_SSID_FILTER: ["ssid"]},
+        sites=SITES,
         clients_response=[CLIENT_1, CLIENT_2, CLIENT_3],
         devices_response=[DEVICE_1, DEVICE_2],
         clients_all_response={},
@@ -267,9 +184,10 @@ async def test_wireless_client_go_wired_issue(hass):
         hass,
         ENTRY_CONFIG,
         options={},
+        sites=SITES,
         clients_response=[client_1_client],
-        devices_response={},
+        devices_response=[],
-        clients_all_response={},
+        clients_all_response=[],
     )
     assert len(hass.states.async_all()) == 3
 
@@ -316,14 +234,14 @@ async def test_restoring_client(hass):
     registry.async_get_or_create(
         device_tracker.DOMAIN,
         unifi.DOMAIN,
-        "{}-mock-site".format(CLIENT_1["mac"]),
+        "{}-site_id".format(CLIENT_1["mac"]),
         suggested_object_id=CLIENT_1["hostname"],
         config_entry=config_entry,
     )
     registry.async_get_or_create(
         device_tracker.DOMAIN,
         unifi.DOMAIN,
-        "{}-mock-site".format(CLIENT_2["mac"]),
+        "{}-site_id".format(CLIENT_2["mac"]),
         suggested_object_id=CLIENT_2["hostname"],
         config_entry=config_entry,
     )
@@ -332,8 +250,9 @@ async def test_restoring_client(hass):
         hass,
         ENTRY_CONFIG,
         options={unifi.CONF_BLOCK_CLIENT: True},
+        sites=SITES,
         clients_response=[CLIENT_2],
-        devices_response={},
+        devices_response=[],
         clients_all_response=[CLIENT_1],
     )
     assert len(hass.states.async_all()) == 4
@@ -348,9 +267,10 @@ async def test_dont_track_clients(hass):
         hass,
         ENTRY_CONFIG,
         options={unifi.controller.CONF_TRACK_CLIENTS: False},
+        sites=SITES,
         clients_response=[CLIENT_1],
         devices_response=[DEVICE_1],
-        clients_all_response={},
+        clients_all_response=[],
     )
     assert len(hass.states.async_all()) == 3
 
@@ -368,9 +288,10 @@ async def test_dont_track_devices(hass):
         hass,
         ENTRY_CONFIG,
         options={unifi.controller.CONF_TRACK_DEVICES: False},
+        sites=SITES,
         clients_response=[CLIENT_1],
         devices_response=[DEVICE_1],
-        clients_all_response={},
+        clients_all_response=[],
     )
     assert len(hass.states.async_all()) == 3
 
@@ -388,9 +309,10 @@ async def test_dont_track_wired_clients(hass):
         hass,
         ENTRY_CONFIG,
         options={unifi.controller.CONF_TRACK_WIRED_CLIENTS: False},
+        sites=SITES,
         clients_response=[CLIENT_1, CLIENT_2],
-        devices_response={},
+        devices_response=[],
-        clients_all_response={},
+        clients_all_response=[],
     )
     assert len(hass.states.async_all()) == 3
 

tests/components/unifi/test_init.py

@@ -4,11 +4,7 @@ from unittest.mock import Mock, patch
 from homeassistant.components import unifi
 from homeassistant.components.unifi import config_flow
 from homeassistant.setup import async_setup_component
-from homeassistant.components.unifi.const import (
+from homeassistant.components.unifi.const import CONF_CONTROLLER, CONF_SITE_ID
-    CONF_CONTROLLER,
-    CONF_SITE_ID,
-    CONTROLLER_ID as CONF_CONTROLLER_ID,
-)
 from homeassistant.const import (
     CONF_HOST,
     CONF_PASSWORD,
@@ -117,8 +113,7 @@ async def test_controller_fail_setup(hass):
         mock_cntrlr.return_value.async_setup.return_value = mock_coro(False)
         assert await unifi.async_setup_entry(hass, entry) is False
 
-    controller_id = CONF_CONTROLLER_ID.format(host="0.0.0.0", site="default")
+    assert hass.data[unifi.DOMAIN] == {}
-    assert controller_id in hass.data[unifi.DOMAIN]
 
 
 async def test_controller_no_mac(hass):

tests/components/unifi/test_sensor.py

@@ -1,29 +1,13 @@
 """UniFi sensor platform tests."""
-from collections import deque
 from copy import deepcopy
 
-from asynctest import patch
-
-from homeassistant import config_entries
 from homeassistant.components import unifi
-from homeassistant.components.unifi.const import (
-    CONF_CONTROLLER,
-    CONF_SITE_ID,
-    CONTROLLER_ID as CONF_CONTROLLER_ID,
-    UNIFI_CONFIG,
-    UNIFI_WIRELESS_CLIENTS,
-)
 from homeassistant.setup import async_setup_component
-from homeassistant.const import (
-    CONF_HOST,
-    CONF_PASSWORD,
-    CONF_PORT,
-    CONF_USERNAME,
-    CONF_VERIFY_SSL,
-)
 
 import homeassistant.components.sensor as sensor
 
+from .test_controller import ENTRY_CONFIG, SITES, setup_unifi_integration
+
 CLIENTS = [
     {
         "hostname": "Wired client hostname",
@@ -53,85 +37,6 @@ CLIENTS = [
     },
 ]
 
-CONTROLLER_DATA = {
-    CONF_HOST: "mock-host",
-    CONF_USERNAME: "mock-user",
-    CONF_PASSWORD: "mock-pswd",
-    CONF_PORT: 1234,
-    CONF_SITE_ID: "mock-site",
-    CONF_VERIFY_SSL: False,
-}
-
-ENTRY_CONFIG = {CONF_CONTROLLER: CONTROLLER_DATA}
-
-CONTROLLER_ID = CONF_CONTROLLER_ID.format(host="mock-host", site="mock-site")
-
-SITES = {"Site name": {"desc": "Site name", "name": "mock-site", "role": "admin"}}
-
-
-async def setup_unifi_integration(
-    hass,
-    config,
-    options,
-    sites,
-    clients_response,
-    devices_response,
-    clients_all_response,
-):
-    """Create the UniFi controller."""
-    hass.data[UNIFI_CONFIG] = []
-    hass.data[UNIFI_WIRELESS_CLIENTS] = unifi.UnifiWirelessClients(hass)
-    config_entry = config_entries.ConfigEntry(
-        version=1,
-        domain=unifi.DOMAIN,
-        title="Mock Title",
-        data=config,
-        source="test",
-        connection_class=config_entries.CONN_CLASS_LOCAL_POLL,
-        system_options={},
-        options=options,
-        entry_id=1,
-    )
-
-    mock_client_responses = deque()
-    mock_client_responses.append(clients_response)
-
-    mock_device_responses = deque()
-    mock_device_responses.append(devices_response)
-
-    mock_client_all_responses = deque()
-    mock_client_all_responses.append(clients_all_response)
-
-    mock_requests = []
-
-    async def mock_request(self, method, path, json=None):
-        mock_requests.append({"method": method, "path": path, "json": json})
-
-        if path == "s/{site}/stat/sta" and mock_client_responses:
-            return mock_client_responses.popleft()
-        if path == "s/{site}/stat/device" and mock_device_responses:
-            return mock_device_responses.popleft()
-        if path == "s/{site}/rest/user" and mock_client_all_responses:
-            return mock_client_all_responses.popleft()
-        return {}
-
-    with patch("aiounifi.Controller.login", return_value=True), patch(
-        "aiounifi.Controller.sites", return_value=sites
-    ), patch("aiounifi.Controller.request", new=mock_request):
-        await unifi.async_setup_entry(hass, config_entry)
-    await hass.async_block_till_done()
-    hass.config_entries._entries.append(config_entry)
-
-    controller_id = unifi.get_controller_id_from_config_entry(config_entry)
-    controller = hass.data[unifi.DOMAIN][controller_id]
-
-    controller.mock_client_responses = mock_client_responses
-    controller.mock_device_responses = mock_device_responses
-    controller.mock_client_all_responses = mock_client_all_responses
-    controller.mock_requests = mock_requests
-
-    return controller
-
 
 async def test_platform_manually_configured(hass):
     """Test that we do not discover anything or try to set up a controller."""
@@ -160,7 +65,7 @@ async def test_no_clients(hass):
     assert len(hass.states.async_all()) == 2
 
 
-async def test_switches(hass):
+async def test_sensors(hass):
     """Test the update_items function with some clients."""
     controller = await setup_unifi_integration(
         hass,

tests/components/unifi/test_switch.py

@@ -1,32 +1,20 @@
 """UniFi POE control platform tests."""
-from collections import deque
 from copy import deepcopy
 
-from asynctest import Mock, patch
-
-import aiounifi
-
 from homeassistant import config_entries
 from homeassistant.components import unifi
-from homeassistant.components.unifi.const import (
-    CONF_CONTROLLER,
-    CONF_SITE_ID,
-    CONTROLLER_ID as CONF_CONTROLLER_ID,
-    UNIFI_CONFIG,
-    UNIFI_WIRELESS_CLIENTS,
-)
 from homeassistant.helpers import entity_registry
 from homeassistant.setup import async_setup_component
-from homeassistant.const import (
-    CONF_HOST,
-    CONF_PASSWORD,
-    CONF_PORT,
-    CONF_USERNAME,
-    CONF_VERIFY_SSL,
-)
 
 import homeassistant.components.switch as switch
 
+from .test_controller import (
+    CONTROLLER_HOST,
+    ENTRY_CONFIG,
+    SITES,
+    setup_unifi_integration,
+)
+
 CLIENT_1 = {
     "hostname": "client_1",
     "ip": "10.0.0.1",
@@ -79,19 +67,6 @@ CLIENT_4 = {
     "wired-rx_bytes": 1234000000,
     "wired-tx_bytes": 5678000000,
 }
-CLOUDKEY = {
-    "hostname": "client_1",
-    "ip": "mock-host",
-    "is_wired": True,
-    "last_seen": 1562600145,
-    "mac": "10:00:00:00:00:01",
-    "name": "Cloud key",
-    "oui": "Producer",
-    "sw_mac": "00:00:00:00:01:01",
-    "sw_port": 1,
-    "wired-rx_bytes": 1234000000,
-    "wired-tx_bytes": 5678000000,
-}
 POE_SWITCH_CLIENTS = [
     {
         "hostname": "client_1",
@@ -211,85 +186,6 @@ UNBLOCKED = {
     "oui": "Producer",
 }
 
-CONTROLLER_DATA = {
-    CONF_HOST: "mock-host",
-    CONF_USERNAME: "mock-user",
-    CONF_PASSWORD: "mock-pswd",
-    CONF_PORT: 1234,
-    CONF_SITE_ID: "mock-site",
-    CONF_VERIFY_SSL: False,
-}
-
-ENTRY_CONFIG = {CONF_CONTROLLER: CONTROLLER_DATA}
-
-CONTROLLER_ID = CONF_CONTROLLER_ID.format(host="mock-host", site="mock-site")
-
-SITES = {"Site name": {"desc": "Site name", "name": "mock-site", "role": "admin"}}
-
-
-async def setup_unifi_integration(
-    hass,
-    config,
-    options,
-    sites,
-    clients_response,
-    devices_response,
-    clients_all_response,
-):
-    """Create the UniFi controller."""
-    hass.data[UNIFI_CONFIG] = []
-    hass.data[UNIFI_WIRELESS_CLIENTS] = unifi.UnifiWirelessClients(hass)
-    config_entry = config_entries.ConfigEntry(
-        version=1,
-        domain=unifi.DOMAIN,
-        title="Mock Title",
-        data=config,
-        source="test",
-        connection_class=config_entries.CONN_CLASS_LOCAL_POLL,
-        system_options={},
-        options=options,
-        entry_id=1,
-    )
-
-    mock_client_responses = deque()
-    mock_client_responses.append(clients_response)
-
-    mock_device_responses = deque()
-    mock_device_responses.append(devices_response)
-
-    mock_client_all_responses = deque()
-    mock_client_all_responses.append(clients_all_response)
-
-    mock_requests = []
-
-    async def mock_request(self, method, path, json=None):
-        mock_requests.append({"method": method, "path": path, "json": json})
-
-        if path == "s/{site}/stat/sta" and mock_client_responses:
-            return mock_client_responses.popleft()
-        if path == "s/{site}/stat/device" and mock_device_responses:
-            return mock_device_responses.popleft()
-        if path == "s/{site}/rest/user" and mock_client_all_responses:
-            return mock_client_all_responses.popleft()
-        return {}
-
-    with patch("aiounifi.Controller.login", return_value=True), patch(
-        "aiounifi.Controller.sites", return_value=sites
-    ), patch("aiounifi.Controller.request", new=mock_request):
-        await unifi.async_setup_entry(hass, config_entry)
-    await hass.async_block_till_done()
-    hass.config_entries._entries.append(config_entry)
-
-    controller_id = unifi.get_controller_id_from_config_entry(config_entry)
-    controller = hass.data[unifi.DOMAIN][controller_id]
-
-    controller.mock_client_responses = mock_client_responses
-    controller.mock_device_responses = mock_device_responses
-    controller.mock_client_all_responses = mock_client_all_responses
-    controller.mock_requests = mock_requests
-
-    return controller
-
 
 async def test_platform_manually_configured(hass):
     """Test that we do not discover anything or try to set up a controller."""
@@ -331,7 +227,7 @@ async def test_controller_not_client(hass):
             unifi.const.CONF_TRACK_DEVICES: False,
         },
         sites=SITES,
-        clients_response=[CLOUDKEY],
+        clients_response=[CONTROLLER_HOST],
         devices_response=[DEVICE_1],
         clients_all_response=[],
     )
@@ -402,7 +298,51 @@ async def test_switches(hass):
     assert unblocked.state == "on"
 
 
-async def test_new_client_discovered(hass):
+async def test_new_client_discovered_on_block_control(hass):
+    """Test if 2nd update has a new client."""
+    controller = await setup_unifi_integration(
+        hass,
+        ENTRY_CONFIG,
+        options={
+            unifi.CONF_BLOCK_CLIENT: [BLOCKED["mac"]],
+            unifi.const.CONF_TRACK_CLIENTS: False,
+            unifi.const.CONF_TRACK_DEVICES: False,
+        },
+        sites=SITES,
+        clients_response=[],
+        devices_response=[],
+        clients_all_response=[BLOCKED],
+    )
+
+    assert len(controller.mock_requests) == 3
+    assert len(hass.states.async_all()) == 4
+
+    controller.mock_client_all_responses.append([BLOCKED])
+
+    # Calling a service will trigger the updates to run
+    await hass.services.async_call(
+        "switch", "turn_off", {"entity_id": "switch.block_client_1"}, blocking=True
+    )
+    assert len(controller.mock_requests) == 7
+    assert len(hass.states.async_all()) == 4
+    assert controller.mock_requests[3] == {
+        "json": {"mac": "00:00:00:00:01:01", "cmd": "block-sta"},
+        "method": "post",
+        "path": "s/{site}/cmd/stamgr/",
+    }
+
+    await hass.services.async_call(
+        "switch", "turn_on", {"entity_id": "switch.block_client_1"}, blocking=True
+    )
+    assert len(controller.mock_requests) == 11
+    assert controller.mock_requests[7] == {
+        "json": {"mac": "00:00:00:00:01:01", "cmd": "unblock-sta"},
+        "method": "post",
+        "path": "s/{site}/cmd/stamgr/",
+    }
+
+
+async def test_new_client_discovered_on_poe_control(hass):
     """Test if 2nd update has a new client."""
     controller = await setup_unifi_integration(
         hass,
@@ -530,59 +470,3 @@ async def test_restoring_client(hass):
 
     device_1 = hass.states.get("switch.client_1")
     assert device_1 is not None
-
-
-async def test_failed_update_failed_login(hass):
-    """Running update can handle a failed login."""
-    controller = await setup_unifi_integration(
-        hass,
-        ENTRY_CONFIG,
-        options={
-            unifi.CONF_BLOCK_CLIENT: ["random mac"],
-            unifi.const.CONF_TRACK_CLIENTS: False,
-            unifi.const.CONF_TRACK_DEVICES: False,
-        },
-        sites=SITES,
-        clients_response=[],
-        devices_response=[],
-        clients_all_response=[],
-    )
-
-    assert len(controller.mock_requests) == 3
-    assert len(hass.states.async_all()) == 2
-
-    with patch.object(
-        controller.api.clients, "update", side_effect=aiounifi.LoginRequired
-    ), patch.object(controller.api, "login", side_effect=aiounifi.AiounifiException):
-        await controller.async_update()
-    await hass.async_block_till_done()
-
-    assert controller.available is False
-
-
-async def test_failed_update_successful_login(hass):
-    """Running update can login when requested."""
-    controller = await setup_unifi_integration(
-        hass,
-        ENTRY_CONFIG,
-        options={
-            unifi.CONF_BLOCK_CLIENT: ["random mac"],
-            unifi.const.CONF_TRACK_CLIENTS: False,
-            unifi.const.CONF_TRACK_DEVICES: False,
-        },
-        sites=SITES,
-        clients_response=[],
-        devices_response=[],
-        clients_all_response=[],
-    )
-
-    assert len(controller.mock_requests) == 3
-    assert len(hass.states.async_all()) == 2
-
-    with patch.object(
-        controller.api.clients, "update", side_effect=aiounifi.LoginRequired
-    ), patch.object(controller.api, "login", return_value=Mock(True)):
-        await controller.async_update()
-    await hass.async_block_till_done()
-
-    assert controller.available is True