Increase test coverage of UniFi integration (#46347)

* Increase coverage of init

* Increase coverage of config_flow

* Improve coverage of controller

* Minor improvement to switch test

* Fix review comment

* Mock websocket class

* Replace the rest of the old websocket event tests

* Improve websocket fixture for cleaner tests

* Fix typing

* Improve connection state signalling based on Martins feedback

* Improve tests of reconnection_mechanisms based on Martins review comments

* Fix unload entry

* Fix isort issue after rebase

* Fix martins comment on not using caplog

* Fix wireless clients test

* Fix martins comments on wireless clients test

homeassistant/components/unifi/config_flow.py

@@ -123,8 +123,9 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
 
                 return await self.async_step_site()
 
-        host = self.config.get(CONF_HOST)
+        if not (host := self.config.get(CONF_HOST, "")) and await async_discover_unifi(
-        if not host and await async_discover_unifi(self.hass):
+            self.hass
+        ):
             host = "unifi"
 
         data = self.reauth_schema or {

homeassistant/components/unifi/controller.py

@@ -415,9 +415,8 @@ class UniFiController:
         If config entry is updated due to reauth flow
         the entry might already have been reset and thus is not available.
         """
-        if config_entry.entry_id not in hass.data[UNIFI_DOMAIN]:
+        if not (controller := hass.data[UNIFI_DOMAIN].get(config_entry.entry_id)):
             return
-        controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
         controller.load_config_entry_options()
         async_dispatcher_send(hass, controller.signal_options_update)
 

tests/components/unifi/conftest.py

@@ -1,9 +1,30 @@
 """Fixtures for UniFi methods."""
+from typing import Optional
 from unittest.mock import patch
 
+from aiounifi.websocket import SIGNAL_CONNECTION_STATE, SIGNAL_DATA
 import pytest
 
 
+@pytest.fixture(autouse=True)
+def mock_unifi_websocket():
+    """No real websocket allowed."""
+    with patch("aiounifi.controller.WSClient") as mock:
+
+        def make_websocket_call(data: Optional[dict] = None, state: str = ""):
+            """Generate a websocket call."""
+            if data:
+                mock.return_value.data = data
+                mock.call_args[1]["callback"](SIGNAL_DATA)
+            elif state:
+                mock.return_value.state = state
+                mock.call_args[1]["callback"](SIGNAL_CONNECTION_STATE)
+            else:
+                raise NotImplementedError
+
+        yield make_websocket_call
+
+
 @pytest.fixture(autouse=True)
 def mock_discovery():
     """No real network traffic allowed."""

tests/components/unifi/test_config_flow.py

@@ -1,9 +1,12 @@
 """Test UniFi config flow."""
+
+import socket
 from unittest.mock import patch
 
 import aiounifi
 
 from homeassistant import config_entries, data_entry_flow, setup
+from homeassistant.components.unifi.config_flow import async_discover_unifi
 from homeassistant.components.unifi.const import (
     CONF_ALLOW_BANDWIDTH_SENSORS,
     CONF_ALLOW_UPTIME_SENSORS,
@@ -151,6 +154,23 @@ async def test_flow_works(hass, aioclient_mock, mock_discovery):
     }
 
 
+async def test_flow_works_negative_discovery(hass, aioclient_mock, mock_discovery):
+    """Test config flow with a negative outcome of async_discovery_unifi."""
+    result = await hass.config_entries.flow.async_init(
+        UNIFI_DOMAIN, context={"source": "user"}
+    )
+
+    assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
+    assert result["step_id"] == "user"
+    assert result["data_schema"]({CONF_USERNAME: "", CONF_PASSWORD: ""}) == {
+        CONF_HOST: "",
+        CONF_USERNAME: "",
+        CONF_PASSWORD: "",
+        CONF_PORT: 443,
+        CONF_VERIFY_SSL: False,
+    }
+
+
 async def test_flow_multiple_sites(hass, aioclient_mock):
     """Test config flow works when finding multiple sites."""
     result = await hass.config_entries.flow.async_init(
@@ -617,3 +637,15 @@ async def test_form_ssdp_gets_form_with_ignored_entry(hass):
         "host": "1.2.3.4",
         "site": "default",
     }
+
+
+async def test_discover_unifi_positive(hass):
+    """Verify positive run of UniFi discovery."""
+    with patch("socket.gethostbyname", return_value=True):
+        assert await async_discover_unifi(hass)
+
+
+async def test_discover_unifi_negative(hass):
+    """Verify negative run of UniFi discovery."""
+    with patch("socket.gethostbyname", side_effect=socket.gaierror):
+        assert await async_discover_unifi(hass) is None

tests/components/unifi/test_controller.py

@@ -1,10 +1,12 @@
 """Test UniFi Controller."""
 
+import asyncio
 from copy import deepcopy
 from datetime import timedelta
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import aiounifi
+from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING
 import pytest
 
 from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
@@ -13,6 +15,8 @@ from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.components.unifi.const import (
     CONF_CONTROLLER,
     CONF_SITE_ID,
+    CONF_TRACK_CLIENTS,
+    CONF_TRACK_DEVICES,
     DEFAULT_ALLOW_BANDWIDTH_SENSORS,
     DEFAULT_ALLOW_UPTIME_SENSORS,
     DEFAULT_DETECTION_TIME,
@@ -22,7 +26,11 @@ from homeassistant.components.unifi.const import (
     DOMAIN as UNIFI_DOMAIN,
     UNIFI_WIRELESS_CLIENTS,
 )
-from homeassistant.components.unifi.controller import PLATFORMS, get_controller
+from homeassistant.components.unifi.controller import (
+    PLATFORMS,
+    RETRY_TIMER,
+    get_controller,
+)
 from homeassistant.components.unifi.errors import AuthenticationRequired, CannotConnect
 from homeassistant.const import (
     CONF_HOST,
@@ -32,10 +40,13 @@ from homeassistant.const import (
     CONF_VERIFY_SSL,
     CONTENT_TYPE_JSON,
 )
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.setup import async_setup_component
+import homeassistant.util.dt as dt_util
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, async_fire_time_changed
 
+DEFAULT_CONFIG_ENTRY_ID = 1
 DEFAULT_HOST = "1.2.3.4"
 DEFAULT_SITE = "site_id"
 
@@ -154,6 +165,7 @@ async def setup_unifi_integration(
     wlans_response=None,
     known_wireless_clients=None,
     controllers=None,
+    unique_id="1",
 ):
     """Create the UniFi controller."""
     assert await async_setup_component(hass, UNIFI_DOMAIN, {})
@@ -162,8 +174,8 @@ async def setup_unifi_integration(
         domain=UNIFI_DOMAIN,
         data=deepcopy(config),
         options=deepcopy(options),
-        entry_id=1,
+        unique_id=unique_id,
-        unique_id="1",
+        entry_id=DEFAULT_CONFIG_ENTRY_ID,
         version=1,
     )
     config_entry.add_to_hass(hass)
@@ -188,8 +200,7 @@ async def setup_unifi_integration(
             wlans_response=wlans_response,
         )
 
-    with patch.object(aiounifi.websocket.WSClient, "start", return_value=True):
+    await hass.config_entries.async_setup(config_entry.entry_id)
-        await hass.config_entries.async_setup(config_entry.entry_id)
     await hass.async_block_till_done()
 
     if config_entry.entry_id not in hass.data[UNIFI_DOMAIN]:
@@ -276,6 +287,27 @@ async def test_controller_unknown_error(hass):
     assert hass.data[UNIFI_DOMAIN] == {}
 
 
+async def test_config_entry_updated(hass, aioclient_mock):
+    """Calling reset when the entry has been setup."""
+    config_entry = await setup_unifi_integration(hass, aioclient_mock)
+    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+
+    event_call = Mock()
+    unsub = async_dispatcher_connect(hass, controller.signal_options_update, event_call)
+
+    hass.config_entries.async_update_entry(
+        config_entry, options={CONF_TRACK_CLIENTS: False, CONF_TRACK_DEVICES: False}
+    )
+    await hass.async_block_till_done()
+
+    assert config_entry.options[CONF_TRACK_CLIENTS] is False
+    assert config_entry.options[CONF_TRACK_DEVICES] is False
+
+    event_call.assert_called_once()
+
+    unsub()
+
+
 async def test_reset_after_successful_setup(hass, aioclient_mock):
     """Calling reset when the entry has been setup."""
     config_entry = await setup_unifi_integration(hass, aioclient_mock)
@@ -290,33 +322,126 @@ async def test_reset_after_successful_setup(hass, aioclient_mock):
     assert len(controller.listeners) == 0
 
 
-async def test_wireless_client_event_calls_update_wireless_devices(
+async def test_reset_fails(hass, aioclient_mock):
-    hass, aioclient_mock
+    """Calling reset when the entry has been setup can return false."""
-):
-    """Call update_wireless_devices method when receiving wireless client event."""
     config_entry = await setup_unifi_integration(hass, aioclient_mock)
     controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
 
+    with patch(
+        "homeassistant.config_entries.ConfigEntries.async_forward_entry_unload",
+        return_value=False,
+    ):
+        result = await controller.async_reset()
+        await hass.async_block_till_done()
+
+    assert result is False
+
+
+async def test_connection_state_signalling(hass, aioclient_mock, mock_unifi_websocket):
+    """Verify connection statesignalling and connection state are working."""
+    client = {
+        "hostname": "client",
+        "ip": "10.0.0.1",
+        "is_wired": True,
+        "last_seen": dt_util.as_timestamp(dt_util.utcnow()),
+        "mac": "00:00:00:00:00:01",
+    }
+    await setup_unifi_integration(hass, aioclient_mock, clients_response=[client])
+
+    # Controller is connected
+    assert hass.states.get("device_tracker.client").state == "home"
+
+    mock_unifi_websocket(state=STATE_DISCONNECTED)
+    await hass.async_block_till_done()
+
+    # Controller is disconnected
+    assert hass.states.get("device_tracker.client").state == "unavailable"
+
+    mock_unifi_websocket(state=STATE_RUNNING)
+    await hass.async_block_till_done()
+
+    # Controller is once again connected
+    assert hass.states.get("device_tracker.client").state == "home"
+
+
+async def test_wireless_client_event_calls_update_wireless_devices(
+    hass, aioclient_mock, mock_unifi_websocket
+):
+    """Call update_wireless_devices method when receiving wireless client event."""
+    await setup_unifi_integration(hass, aioclient_mock)
+
     with patch(
         "homeassistant.components.unifi.controller.UniFiController.update_wireless_clients",
         return_value=None,
     ) as wireless_clients_mock:
-        controller.api.websocket._data = {
+        mock_unifi_websocket(
-            "meta": {"rc": "ok", "message": "events"},
+            data={
-            "data": [
+                "meta": {"rc": "ok", "message": "events"},
-                {
+                "data": [
-                    "datetime": "2020-01-20T19:37:04Z",
+                    {
-                    "key": aiounifi.events.WIRELESS_CLIENT_CONNECTED,
+                        "datetime": "2020-01-20T19:37:04Z",
-                    "msg": "User[11:22:33:44:55:66] has connected to WLAN",
+                        "key": aiounifi.events.WIRELESS_CLIENT_CONNECTED,
-                    "time": 1579549024893,
+                        "msg": "User[11:22:33:44:55:66] has connected to WLAN",
-                }
+                        "time": 1579549024893,
-            ],
+                    }
-        }
+                ],
-        controller.api.session_handler("data")
+            },
+        )
 
         assert wireless_clients_mock.assert_called_once
 
 
+async def test_reconnect_mechanism(hass, aioclient_mock, mock_unifi_websocket):
+    """Verify reconnect prints only on first reconnection try."""
+    await setup_unifi_integration(hass, aioclient_mock)
+
+    aioclient_mock.clear_requests()
+    aioclient_mock.post(f"https://{DEFAULT_HOST}:1234/api/login", status=502)
+
+    mock_unifi_websocket(state=STATE_DISCONNECTED)
+    await hass.async_block_till_done()
+
+    assert aioclient_mock.call_count == 0
+
+    new_time = dt_util.utcnow() + timedelta(seconds=RETRY_TIMER)
+    async_fire_time_changed(hass, new_time)
+    await hass.async_block_till_done()
+
+    assert aioclient_mock.call_count == 1
+
+    new_time = dt_util.utcnow() + timedelta(seconds=RETRY_TIMER)
+    async_fire_time_changed(hass, new_time)
+    await hass.async_block_till_done()
+
+    assert aioclient_mock.call_count == 2
+
+
+@pytest.mark.parametrize(
+    "exception",
+    [
+        asyncio.TimeoutError,
+        aiounifi.BadGateway,
+        aiounifi.ServiceUnavailable,
+        aiounifi.AiounifiException,
+    ],
+)
+async def test_reconnect_mechanism_exceptions(
+    hass, aioclient_mock, mock_unifi_websocket, exception
+):
+    """Verify async_reconnect calls expected methods."""
+    await setup_unifi_integration(hass, aioclient_mock)
+
+    with patch("aiounifi.Controller.login", side_effect=exception), patch(
+        "homeassistant.components.unifi.controller.UniFiController.reconnect"
+    ) as mock_reconnect:
+        mock_unifi_websocket(state=STATE_DISCONNECTED)
+        await hass.async_block_till_done()
+
+        new_time = dt_util.utcnow() + timedelta(seconds=RETRY_TIMER)
+        async_fire_time_changed(hass, new_time)
+        mock_reconnect.assert_called_once()
+
+
 async def test_get_controller(hass):
     """Successful call."""
     with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(

tests/components/unifi/test_device_tracker.py

@@ -8,9 +8,8 @@ from aiounifi.controller import (
     MESSAGE_CLIENT_REMOVED,
     MESSAGE_DEVICE,
     MESSAGE_EVENT,
-    SIGNAL_CONNECTION_STATE,
 )
-from aiounifi.websocket import SIGNAL_DATA, STATE_DISCONNECTED, STATE_RUNNING
+from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING
 
 from homeassistant import config_entries
 from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
@@ -157,7 +156,7 @@ async def test_no_clients(hass, aioclient_mock):
     assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 0
 
 
-async def test_tracked_wireless_clients(hass, aioclient_mock):
+async def test_tracked_wireless_clients(hass, aioclient_mock, mock_unifi_websocket):
     """Test the update_items function with some clients."""
     config_entry = await setup_unifi_integration(
         hass, aioclient_mock, clients_response=[CLIENT_1]
@@ -171,11 +170,12 @@ async def test_tracked_wireless_clients(hass, aioclient_mock):
 
     # State change signalling works without events
     client_1_copy = copy(CLIENT_1)
-    controller.api.websocket._data = {
+    mock_unifi_websocket(
-        "meta": {"message": MESSAGE_CLIENT},
+        data={
-        "data": [client_1_copy],
+            "meta": {"message": MESSAGE_CLIENT},
-    }
+            "data": [client_1_copy],
-    controller.api.session_handler(SIGNAL_DATA)
+        }
+    )
     await hass.async_block_till_done()
 
     client_1 = hass.states.get("device_tracker.client_1")
@@ -186,11 +186,13 @@ async def test_tracked_wireless_clients(hass, aioclient_mock):
     assert client_1.attributes["host_name"] == "client_1"
 
     # State change signalling works with events
-    controller.api.websocket._data = {
+
-        "meta": {"message": MESSAGE_EVENT},
+    mock_unifi_websocket(
-        "data": [EVENT_CLIENT_1_WIRELESS_DISCONNECTED],
+        data={
-    }
+            "meta": {"message": MESSAGE_EVENT},
-    controller.api.session_handler(SIGNAL_DATA)
+            "data": [EVENT_CLIENT_1_WIRELESS_DISCONNECTED],
+        }
+    )
     await hass.async_block_till_done()
 
     client_1 = hass.states.get("device_tracker.client_1")
@@ -204,30 +206,30 @@ async def test_tracked_wireless_clients(hass, aioclient_mock):
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "not_home"
 
-    controller.api.websocket._data = {
+    mock_unifi_websocket(
-        "meta": {"message": MESSAGE_EVENT},
+        data={
-        "data": [EVENT_CLIENT_1_WIRELESS_CONNECTED],
+            "meta": {"message": MESSAGE_EVENT},
-    }
+            "data": [EVENT_CLIENT_1_WIRELESS_CONNECTED],
-    controller.api.session_handler(SIGNAL_DATA)
+        }
+    )
     await hass.async_block_till_done()
 
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "home"
 
 
-async def test_tracked_clients(hass, aioclient_mock):
+async def test_tracked_clients(hass, aioclient_mock, mock_unifi_websocket):
     """Test the update_items function with some clients."""
     client_4_copy = copy(CLIENT_4)
     client_4_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
 
-    config_entry = await setup_unifi_integration(
+    await setup_unifi_integration(
         hass,
         aioclient_mock,
         options={CONF_SSID_FILTER: ["ssid"]},
         clients_response=[CLIENT_1, CLIENT_2, CLIENT_3, CLIENT_5, client_4_copy],
         known_wireless_clients=(CLIENT_4["mac"],),
     )
-    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
     assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 4
 
     client_1 = hass.states.get("device_tracker.client_1")
@@ -254,22 +256,26 @@ async def test_tracked_clients(hass, aioclient_mock):
 
     # State change signalling works
     client_1_copy = copy(CLIENT_1)
-    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_copy]}
+    mock_unifi_websocket(
-    controller.api.message_handler(event)
+        data={
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": [client_1_copy],
+        }
+    )
     await hass.async_block_till_done()
 
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "home"
 
 
-async def test_tracked_devices(hass, aioclient_mock):
+async def test_tracked_devices(hass, aioclient_mock, mock_unifi_websocket):
     """Test the update_items function with some devices."""
-    config_entry = await setup_unifi_integration(
+    await setup_unifi_integration(
         hass,
         aioclient_mock,
         devices_response=[DEVICE_1, DEVICE_2],
     )
-    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+
     assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
 
     device_1 = hass.states.get("device_tracker.device_1")
@@ -283,12 +289,20 @@ async def test_tracked_devices(hass, aioclient_mock):
     # State change signalling work
     device_1_copy = copy(DEVICE_1)
     device_1_copy["next_interval"] = 20
-    event = {"meta": {"message": MESSAGE_DEVICE}, "data": [device_1_copy]}
+    mock_unifi_websocket(
-    controller.api.message_handler(event)
+        data={
+            "meta": {"message": MESSAGE_DEVICE},
+            "data": [device_1_copy],
+        }
+    )
     device_2_copy = copy(DEVICE_2)
     device_2_copy["next_interval"] = 50
-    event = {"meta": {"message": MESSAGE_DEVICE}, "data": [device_2_copy]}
+    mock_unifi_websocket(
-    controller.api.message_handler(event)
+        data={
+            "meta": {"message": MESSAGE_DEVICE},
+            "data": [device_2_copy],
+        }
+    )
     await hass.async_block_till_done()
 
     device_1 = hass.states.get("device_tracker.device_1")
@@ -309,8 +323,12 @@ async def test_tracked_devices(hass, aioclient_mock):
     # Disabled device is unavailable
     device_1_copy = copy(DEVICE_1)
     device_1_copy["disabled"] = True
-    event = {"meta": {"message": MESSAGE_DEVICE}, "data": [device_1_copy]}
+    mock_unifi_websocket(
-    controller.api.message_handler(event)
+        data={
+            "meta": {"message": MESSAGE_DEVICE},
+            "data": [device_1_copy],
+        }
+    )
     await hass.async_block_till_done()
 
     device_1 = hass.states.get("device_tracker.device_1")
@@ -319,10 +337,18 @@ async def test_tracked_devices(hass, aioclient_mock):
     # Update device registry when device is upgraded
     device_2_copy = copy(DEVICE_2)
     device_2_copy["version"] = EVENT_DEVICE_2_UPGRADED["version_to"]
-    message = {"meta": {"message": MESSAGE_DEVICE}, "data": [device_2_copy]}
+    mock_unifi_websocket(
-    controller.api.message_handler(message)
+        data={
-    event = {"meta": {"message": MESSAGE_EVENT}, "data": [EVENT_DEVICE_2_UPGRADED]}
+            "meta": {"message": MESSAGE_DEVICE},
-    controller.api.message_handler(event)
+            "data": [device_2_copy],
+        }
+    )
+    mock_unifi_websocket(
+        data={
+            "meta": {"message": MESSAGE_EVENT},
+            "data": [EVENT_DEVICE_2_UPGRADED],
+        }
+    )
     await hass.async_block_till_done()
 
     # Verify device registry has been updated
@@ -333,12 +359,12 @@ async def test_tracked_devices(hass, aioclient_mock):
     assert device.sw_version == EVENT_DEVICE_2_UPGRADED["version_to"]
 
 
-async def test_remove_clients(hass, aioclient_mock):
+async def test_remove_clients(hass, aioclient_mock, mock_unifi_websocket):
     """Test the remove_items function with some clients."""
-    config_entry = await setup_unifi_integration(
+    await setup_unifi_integration(
         hass, aioclient_mock, clients_response=[CLIENT_1, CLIENT_2]
     )
-    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+
     assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
 
     client_1 = hass.states.get("device_tracker.client_1")
@@ -347,11 +373,12 @@ async def test_remove_clients(hass, aioclient_mock):
     wired_client = hass.states.get("device_tracker.wired_client")
     assert wired_client is not None
 
-    controller.api.websocket._data = {
+    mock_unifi_websocket(
-        "meta": {"message": MESSAGE_CLIENT_REMOVED},
+        data={
-        "data": [CLIENT_1],
+            "meta": {"message": MESSAGE_CLIENT_REMOVED},
-    }
+            "data": [CLIENT_1],
-    controller.api.session_handler(SIGNAL_DATA)
+        }
+    )
     await hass.async_block_till_done()
     await hass.async_block_till_done()
 
@@ -364,15 +391,15 @@ async def test_remove_clients(hass, aioclient_mock):
     assert wired_client is not None
 
 
-async def test_controller_state_change(hass, aioclient_mock):
+async def test_controller_state_change(hass, aioclient_mock, mock_unifi_websocket):
     """Verify entities state reflect on controller becoming unavailable."""
-    config_entry = await setup_unifi_integration(
+    await setup_unifi_integration(
         hass,
         aioclient_mock,
         clients_response=[CLIENT_1],
         devices_response=[DEVICE_1],
     )
-    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+
     assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
 
     client_1 = hass.states.get("device_tracker.client_1")
@@ -382,9 +409,7 @@ async def test_controller_state_change(hass, aioclient_mock):
     assert device_1.state == "home"
 
     # Controller unavailable
-    controller.async_unifi_signalling_callback(
+    mock_unifi_websocket(state=STATE_DISCONNECTED)
-        SIGNAL_CONNECTION_STATE, STATE_DISCONNECTED
-    )
     await hass.async_block_till_done()
 
     client_1 = hass.states.get("device_tracker.client_1")
@@ -394,7 +419,7 @@ async def test_controller_state_change(hass, aioclient_mock):
     assert device_1.state == STATE_UNAVAILABLE
 
     # Controller available
-    controller.async_unifi_signalling_callback(SIGNAL_CONNECTION_STATE, STATE_RUNNING)
+    mock_unifi_websocket(state=STATE_RUNNING)
     await hass.async_block_till_done()
 
     client_1 = hass.states.get("device_tracker.client_1")
@@ -554,7 +579,7 @@ async def test_option_track_devices(hass, aioclient_mock):
     assert device_1 is not None
 
 
-async def test_option_ssid_filter(hass, aioclient_mock):
+async def test_option_ssid_filter(hass, aioclient_mock, mock_unifi_websocket):
     """Test the SSID filter works.
 
     Client 1 will travel from a supported SSID to an unsupported ssid.
@@ -593,13 +618,21 @@ async def test_option_ssid_filter(hass, aioclient_mock):
     # Roams to SSID outside of filter
     client_1_copy = copy(CLIENT_1)
     client_1_copy["essid"] = "other_ssid"
-    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_copy]}
+    mock_unifi_websocket(
-    controller.api.message_handler(event)
+        data={
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": [client_1_copy],
+        }
+    )
     # Data update while SSID filter is in effect shouldn't create the client
     client_3_copy = copy(CLIENT_3)
     client_3_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
-    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_3_copy]}
+    mock_unifi_websocket(
-    controller.api.message_handler(event)
+        data={
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": [client_3_copy],
+        }
+    )
     await hass.async_block_till_done()
 
     # SSID filter marks client as away
@@ -616,10 +649,19 @@ async def test_option_ssid_filter(hass, aioclient_mock):
         options={CONF_SSID_FILTER: []},
     )
     await hass.async_block_till_done()
-    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_copy]}
+
-    controller.api.message_handler(event)
+    mock_unifi_websocket(
-    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_3_copy]}
+        data={
-    controller.api.message_handler(event)
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": [client_1_copy],
+        }
+    )
+    mock_unifi_websocket(
+        data={
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": [client_3_copy],
+        }
+    )
     await hass.async_block_till_done()
 
     client_1 = hass.states.get("device_tracker.client_1")
@@ -636,16 +678,24 @@ async def test_option_ssid_filter(hass, aioclient_mock):
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "not_home"
 
-    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_3_copy]}
+    mock_unifi_websocket(
-    controller.api.message_handler(event)
+        data={
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": [client_3_copy],
+        }
+    )
     await hass.async_block_till_done()
     # Client won't go away until after next update
     client_3 = hass.states.get("device_tracker.client_3")
     assert client_3.state == "home"
 
     # Trigger update to get client marked as away
-    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [CLIENT_3]}
+    mock_unifi_websocket(
-    controller.api.message_handler(event)
+        data={
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": [client_3_copy],
+        }
+    )
     await hass.async_block_till_done()
 
     new_time = (
@@ -659,7 +709,9 @@ async def test_option_ssid_filter(hass, aioclient_mock):
     assert client_3.state == "not_home"
 
 
-async def test_wireless_client_go_wired_issue(hass, aioclient_mock):
+async def test_wireless_client_go_wired_issue(
+    hass, aioclient_mock, mock_unifi_websocket
+):
     """Test the solution to catch wireless device go wired UniFi issue.
 
     UniFi has a known issue that when a wireless device goes away it sometimes gets marked as wired.
@@ -681,8 +733,12 @@ async def test_wireless_client_go_wired_issue(hass, aioclient_mock):
 
     # Trigger wired bug
     client_1_client["is_wired"] = True
-    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
+    mock_unifi_websocket(
-    controller.api.message_handler(event)
+        data={
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": [client_1_client],
+        }
+    )
     await hass.async_block_till_done()
 
     # Wired bug fix keeps client marked as wireless
@@ -702,8 +758,12 @@ async def test_wireless_client_go_wired_issue(hass, aioclient_mock):
     assert client_1.attributes["is_wired"] is False
 
     # Try to mark client as connected
-    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
+    mock_unifi_websocket(
-    controller.api.message_handler(event)
+        data={
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": [client_1_client],
+        }
+    )
     await hass.async_block_till_done()
 
     # Make sure it don't go online again until wired bug disappears
@@ -713,8 +773,12 @@ async def test_wireless_client_go_wired_issue(hass, aioclient_mock):
 
     # Make client wireless
     client_1_client["is_wired"] = False
-    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
+    mock_unifi_websocket(
-    controller.api.message_handler(event)
+        data={
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": [client_1_client],
+        }
+    )
     await hass.async_block_till_done()
 
     # Client is no longer affected by wired bug and can be marked online
@@ -723,7 +787,7 @@ async def test_wireless_client_go_wired_issue(hass, aioclient_mock):
     assert client_1.attributes["is_wired"] is False
 
 
-async def test_option_ignore_wired_bug(hass, aioclient_mock):
+async def test_option_ignore_wired_bug(hass, aioclient_mock, mock_unifi_websocket):
     """Test option to ignore wired bug."""
     client_1_client = copy(CLIENT_1)
     client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
@@ -745,8 +809,12 @@ async def test_option_ignore_wired_bug(hass, aioclient_mock):
 
     # Trigger wired bug
     client_1_client["is_wired"] = True
-    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
+    mock_unifi_websocket(
-    controller.api.message_handler(event)
+        data={
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": [client_1_client],
+        }
+    )
     await hass.async_block_till_done()
 
     # Wired bug in effect
@@ -766,8 +834,12 @@ async def test_option_ignore_wired_bug(hass, aioclient_mock):
     assert client_1.attributes["is_wired"] is True
 
     # Mark client as connected again
-    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
+    mock_unifi_websocket(
-    controller.api.message_handler(event)
+        data={
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": [client_1_client],
+        }
+    )
     await hass.async_block_till_done()
 
     # Ignoring wired bug allows client to go home again even while affected
@@ -777,8 +849,12 @@ async def test_option_ignore_wired_bug(hass, aioclient_mock):
 
     # Make client wireless
     client_1_client["is_wired"] = False
-    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
+    mock_unifi_websocket(
-    controller.api.message_handler(event)
+        data={
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": [client_1_client],
+        }
+    )
     await hass.async_block_till_done()
 
     # Client is wireless and still connected

tests/components/unifi/test_init.py

@@ -1,14 +1,20 @@
 """Test UniFi setup process."""
-from unittest.mock import AsyncMock, Mock, patch
+from unittest.mock import AsyncMock, patch
 
 from homeassistant.components import unifi
 from homeassistant.components.unifi import async_flatten_entry_data
 from homeassistant.components.unifi.const import CONF_CONTROLLER, DOMAIN as UNIFI_DOMAIN
+from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 from homeassistant.setup import async_setup_component
 
-from .test_controller import CONTROLLER_DATA, ENTRY_CONFIG, setup_unifi_integration
+from .test_controller import (
+    CONTROLLER_DATA,
+    DEFAULT_CONFIG_ENTRY_ID,
+    ENTRY_CONFIG,
+    setup_unifi_integration,
+)
 
-from tests.common import MockConfigEntry, mock_coro
+from tests.common import MockConfigEntry
 
 
 async def test_setup_with_no_config(hass):
@@ -19,7 +25,7 @@ async def test_setup_with_no_config(hass):
 
 async def test_successful_config_entry(hass, aioclient_mock):
     """Test that configured options for a host are loaded via config entry."""
-    await setup_unifi_integration(hass, aioclient_mock)
+    await setup_unifi_integration(hass, aioclient_mock, unique_id=None)
     assert hass.data[UNIFI_DOMAIN]
 
 
@@ -32,29 +38,28 @@ async def test_controller_fail_setup(hass):
     assert hass.data[UNIFI_DOMAIN] == {}
 
 
-async def test_controller_no_mac(hass):
+async def test_controller_mac(hass):
     """Test that configured options for a host are loaded via config entry."""
     entry = MockConfigEntry(
-        domain=UNIFI_DOMAIN,
+        domain=UNIFI_DOMAIN, data=ENTRY_CONFIG, unique_id="1", entry_id=1
-        data=ENTRY_CONFIG,
-        unique_id="1",
-        version=1,
     )
     entry.add_to_hass(hass)
-    mock_registry = Mock()
+
-    with patch(
+    with patch("homeassistant.components.unifi.UniFiController") as mock_controller:
-        "homeassistant.components.unifi.UniFiController"
-    ) as mock_controller, patch(
-        "homeassistant.helpers.device_registry.async_get_registry",
-        return_value=mock_coro(mock_registry),
-    ):
         mock_controller.return_value.async_setup = AsyncMock(return_value=True)
-        mock_controller.return_value.mac = None
+        mock_controller.return_value.mac = "mac1"
         assert await unifi.async_setup_entry(hass, entry) is True
 
     assert len(mock_controller.mock_calls) == 2
 
-    assert len(mock_registry.mock_calls) == 0
+    device_registry = await hass.helpers.device_registry.async_get_registry()
+    device = device_registry.async_get_or_create(
+        config_entry_id=entry.entry_id, connections={(CONNECTION_NETWORK_MAC, "mac1")}
+    )
+    assert device.manufacturer == "Ubiquiti Networks"
+    assert device.model == "UniFi Controller"
+    assert device.name == "UniFi Controller"
+    assert device.sw_version is None
 
 
 async def test_flatten_entry_data(hass):
@@ -73,5 +78,45 @@ async def test_unload_entry(hass, aioclient_mock):
     config_entry = await setup_unifi_integration(hass, aioclient_mock)
     assert hass.data[UNIFI_DOMAIN]
 
-    assert await unifi.async_unload_entry(hass, config_entry)
+    assert await hass.config_entries.async_unload(config_entry.entry_id)
     assert not hass.data[UNIFI_DOMAIN]
+
+
+async def test_wireless_clients(hass, hass_storage, aioclient_mock):
+    """Verify wireless clients class."""
+    hass_storage[unifi.STORAGE_KEY] = {
+        "version": unifi.STORAGE_VERSION,
+        "data": {
+            DEFAULT_CONFIG_ENTRY_ID: {
+                "wireless_devices": ["00:00:00:00:00:00", "00:00:00:00:00:01"]
+            }
+        },
+    }
+
+    client_1 = {
+        "hostname": "client_1",
+        "ip": "10.0.0.1",
+        "is_wired": False,
+        "mac": "00:00:00:00:00:01",
+    }
+    client_2 = {
+        "hostname": "client_2",
+        "ip": "10.0.0.2",
+        "is_wired": False,
+        "mac": "00:00:00:00:00:02",
+    }
+    config_entry = await setup_unifi_integration(
+        hass, aioclient_mock, clients_response=[client_1, client_2]
+    )
+
+    for mac in [
+        "00:00:00:00:00:00",
+        "00:00:00:00:00:01",
+        "00:00:00:00:00:02",
+    ]:
+        assert (
+            mac
+            in hass_storage[unifi.STORAGE_KEY]["data"][config_entry.entry_id][
+                "wireless_devices"
+            ]
+        )

tests/components/unifi/test_sensor.py

@@ -2,7 +2,6 @@
 from copy import deepcopy
 
 from aiounifi.controller import MESSAGE_CLIENT, MESSAGE_CLIENT_REMOVED
-from aiounifi.websocket import SIGNAL_DATA
 
 from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
@@ -63,7 +62,7 @@ async def test_no_clients(hass, aioclient_mock):
     assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 0
 
 
-async def test_sensors(hass, aioclient_mock):
+async def test_sensors(hass, aioclient_mock, mock_unifi_websocket):
     """Test the update_items function with some clients."""
     config_entry = await setup_unifi_integration(
         hass,
@@ -104,8 +103,12 @@ async def test_sensors(hass, aioclient_mock):
     clients[1]["tx_bytes"] = 6789000000
     clients[1]["uptime"] = 1600180860
 
-    event = {"meta": {"message": MESSAGE_CLIENT}, "data": clients}
+    mock_unifi_websocket(
-    controller.api.message_handler(event)
+        data={
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": clients,
+        }
+    )
     await hass.async_block_till_done()
 
     wireless_client_rx = hass.states.get("sensor.wireless_client_name_rx")
@@ -178,9 +181,9 @@ async def test_sensors(hass, aioclient_mock):
     assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 6
 
 
-async def test_remove_sensors(hass, aioclient_mock):
+async def test_remove_sensors(hass, aioclient_mock, mock_unifi_websocket):
     """Test the remove_items function with some clients."""
-    config_entry = await setup_unifi_integration(
+    await setup_unifi_integration(
         hass,
         aioclient_mock,
         options={
@@ -189,7 +192,7 @@ async def test_remove_sensors(hass, aioclient_mock):
         },
         clients_response=CLIENTS,
     )
-    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+
     assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 6
     assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
 
@@ -209,11 +212,12 @@ async def test_remove_sensors(hass, aioclient_mock):
     wireless_client_uptime = hass.states.get("sensor.wireless_client_name_uptime")
     assert wireless_client_uptime is not None
 
-    controller.api.websocket._data = {
+    mock_unifi_websocket(
-        "meta": {"message": MESSAGE_CLIENT_REMOVED},
+        data={
-        "data": [CLIENTS[0]],
+            "meta": {"message": MESSAGE_CLIENT_REMOVED},
-    }
+            "data": [CLIENTS[0]],
-    controller.api.session_handler(SIGNAL_DATA)
+        }
+    )
     await hass.async_block_till_done()
 
     assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 3

tests/components/unifi/test_switch.py

@@ -2,7 +2,6 @@
 from copy import deepcopy
 
 from aiounifi.controller import MESSAGE_CLIENT_REMOVED, MESSAGE_EVENT
-from aiounifi.websocket import SIGNAL_DATA
 
 from homeassistant import config_entries
 from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
@@ -17,6 +16,7 @@ from homeassistant.components.unifi.const import (
 )
 from homeassistant.components.unifi.switch import POE_SWITCH
 from homeassistant.helpers import entity_registry
+from homeassistant.helpers.dispatcher import async_dispatcher_send
 
 from .test_controller import (
     CONTROLLER_HOST,
@@ -370,6 +370,7 @@ async def test_switches(hass, aioclient_mock):
     dpi_switch = hass.states.get("switch.block_media_streaming")
     assert dpi_switch is not None
     assert dpi_switch.state == "on"
+    assert dpi_switch.attributes["icon"] == "mdi:network"
 
     # Block and unblock client
 
@@ -419,17 +420,22 @@ async def test_switches(hass, aioclient_mock):
     assert aioclient_mock.call_count == 14
     assert aioclient_mock.mock_calls[13][2] == {"enabled": True}
 
+    # Make sure no duplicates arise on generic signal update
+    async_dispatcher_send(hass, controller.signal_update)
+    await hass.async_block_till_done()
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 4
 
-async def test_remove_switches(hass, aioclient_mock):
+
+async def test_remove_switches(hass, aioclient_mock, mock_unifi_websocket):
     """Test the update_items function with some clients."""
-    config_entry = await setup_unifi_integration(
+    await setup_unifi_integration(
         hass,
         aioclient_mock,
         options={CONF_BLOCK_CLIENT: [UNBLOCKED["mac"]]},
         clients_response=[CLIENT_1, UNBLOCKED],
         devices_response=[DEVICE_1],
     )
-    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
 
     poe_switch = hass.states.get("switch.poe_client_1")
@@ -438,11 +444,12 @@ async def test_remove_switches(hass, aioclient_mock):
     block_switch = hass.states.get("switch.block_client_2")
     assert block_switch is not None
 
-    controller.api.websocket._data = {
+    mock_unifi_websocket(
-        "meta": {"message": MESSAGE_CLIENT_REMOVED},
+        data={
-        "data": [CLIENT_1, UNBLOCKED],
+            "meta": {"message": MESSAGE_CLIENT_REMOVED},
-    }
+            "data": [CLIENT_1, UNBLOCKED],
-    controller.api.session_handler(SIGNAL_DATA)
+        }
+    )
     await hass.async_block_till_done()
 
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
@@ -454,7 +461,7 @@ async def test_remove_switches(hass, aioclient_mock):
     assert block_switch is None
 
 
-async def test_block_switches(hass, aioclient_mock):
+async def test_block_switches(hass, aioclient_mock, mock_unifi_websocket):
     """Test the update_items function with some clients."""
     config_entry = await setup_unifi_integration(
         hass,
@@ -479,11 +486,12 @@ async def test_block_switches(hass, aioclient_mock):
     assert unblocked is not None
     assert unblocked.state == "on"
 
-    controller.api.websocket._data = {
+    mock_unifi_websocket(
-        "meta": {"message": MESSAGE_EVENT},
+        data={
-        "data": [EVENT_BLOCKED_CLIENT_UNBLOCKED],
+            "meta": {"message": MESSAGE_EVENT},
-    }
+            "data": [EVENT_BLOCKED_CLIENT_UNBLOCKED],
-    controller.api.session_handler(SIGNAL_DATA)
+        }
+    )
     await hass.async_block_till_done()
 
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
@@ -491,11 +499,12 @@ async def test_block_switches(hass, aioclient_mock):
     assert blocked is not None
     assert blocked.state == "on"
 
-    controller.api.websocket._data = {
+    mock_unifi_websocket(
-        "meta": {"message": MESSAGE_EVENT},
+        data={
-        "data": [EVENT_BLOCKED_CLIENT_BLOCKED],
+            "meta": {"message": MESSAGE_EVENT},
-    }
+            "data": [EVENT_BLOCKED_CLIENT_BLOCKED],
-    controller.api.session_handler(SIGNAL_DATA)
+        }
+    )
     await hass.async_block_till_done()
 
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
@@ -526,9 +535,11 @@ async def test_block_switches(hass, aioclient_mock):
     }
 
 
-async def test_new_client_discovered_on_block_control(hass, aioclient_mock):
+async def test_new_client_discovered_on_block_control(
+    hass, aioclient_mock, mock_unifi_websocket
+):
     """Test if 2nd update has a new client."""
-    config_entry = await setup_unifi_integration(
+    await setup_unifi_integration(
         hass,
         aioclient_mock,
         options={
@@ -538,27 +549,28 @@ async def test_new_client_discovered_on_block_control(hass, aioclient_mock):
             CONF_DPI_RESTRICTIONS: False,
         },
     )
-    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
 
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
 
     blocked = hass.states.get("switch.block_client_1")
     assert blocked is None
 
-    controller.api.websocket._data = {
+    mock_unifi_websocket(
-        "meta": {"message": "sta:sync"},
+        data={
-        "data": [BLOCKED],
+            "meta": {"message": "sta:sync"},
-    }
+            "data": [BLOCKED],
-    controller.api.session_handler(SIGNAL_DATA)
+        }
+    )
     await hass.async_block_till_done()
 
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
 
-    controller.api.websocket._data = {
+    mock_unifi_websocket(
-        "meta": {"message": MESSAGE_EVENT},
+        data={
-        "data": [EVENT_BLOCKED_CLIENT_CONNECTED],
+            "meta": {"message": MESSAGE_EVENT},
-    }
+            "data": [EVENT_BLOCKED_CLIENT_CONNECTED],
-    controller.api.session_handler(SIGNAL_DATA)
+        }
+    )
     await hass.async_block_till_done()
 
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
@@ -634,7 +646,9 @@ async def test_option_remove_switches(hass, aioclient_mock):
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
 
 
-async def test_new_client_discovered_on_poe_control(hass, aioclient_mock):
+async def test_new_client_discovered_on_poe_control(
+    hass, aioclient_mock, mock_unifi_websocket
+):
     """Test if 2nd update has a new client."""
     config_entry = await setup_unifi_integration(
         hass,
@@ -647,20 +661,22 @@ async def test_new_client_discovered_on_poe_control(hass, aioclient_mock):
 
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
 
-    controller.api.websocket._data = {
+    mock_unifi_websocket(
-        "meta": {"message": "sta:sync"},
+        data={
-        "data": [CLIENT_2],
+            "meta": {"message": "sta:sync"},
-    }
+            "data": [CLIENT_2],
-    controller.api.session_handler(SIGNAL_DATA)
+        }
+    )
     await hass.async_block_till_done()
 
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
 
-    controller.api.websocket._data = {
+    mock_unifi_websocket(
-        "meta": {"message": MESSAGE_EVENT},
+        data={
-        "data": [EVENT_CLIENT_2_CONNECTED],
+            "meta": {"message": MESSAGE_EVENT},
-    }
+            "data": [EVENT_CLIENT_2_CONNECTED],
-    controller.api.session_handler(SIGNAL_DATA)
+        }
+    )
     await hass.async_block_till_done()
 
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2