From 793929f2ea12f2c04b88bc9be9f82c7a09f55b4a Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Fri, 5 Mar 2021 21:28:41 +0100 Subject: [PATCH] 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 | 5 +- homeassistant/components/unifi/controller.py | 3 +- tests/components/unifi/conftest.py | 21 ++ tests/components/unifi/test_config_flow.py | 32 +++ tests/components/unifi/test_controller.py | 171 +++++++++++-- tests/components/unifi/test_device_tracker.py | 232 ++++++++++++------ tests/components/unifi/test_init.py | 83 +++++-- tests/components/unifi/test_sensor.py | 28 ++- tests/components/unifi/test_switch.py | 104 ++++---- 9 files changed, 499 insertions(+), 180 deletions(-) diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index 5a0a4969f09..6d8c37e8b04 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/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 and await async_discover_unifi(self.hass): + if not (host := self.config.get(CONF_HOST, "")) and await async_discover_unifi( + self.hass + ): host = "unifi" data = self.reauth_schema or { diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 45268a341e1..7c95058e8e4 100644 --- a/homeassistant/components/unifi/controller.py +++ b/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) diff --git a/tests/components/unifi/conftest.py b/tests/components/unifi/conftest.py index b0491a9fa2a..83dc99fdaf8 100644 --- a/tests/components/unifi/conftest.py +++ b/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.""" diff --git a/tests/components/unifi/test_config_flow.py b/tests/components/unifi/test_config_flow.py index a28f5f5f7c5..106c1852414 100644 --- a/tests/components/unifi/test_config_flow.py +++ b/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 diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index 6b40df5857f..50d464d23c0 100644 --- a/tests/components/unifi/test_controller.py +++ b/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="1", + unique_id=unique_id, + 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( - hass, aioclient_mock -): - """Call update_wireless_devices method when receiving wireless client event.""" +async def test_reset_fails(hass, aioclient_mock): + """Calling reset when the entry has been setup can return false.""" 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 = { - "meta": {"rc": "ok", "message": "events"}, - "data": [ - { - "datetime": "2020-01-20T19:37:04Z", - "key": aiounifi.events.WIRELESS_CLIENT_CONNECTED, - "msg": "User[11:22:33:44:55:66] has connected to WLAN", - "time": 1579549024893, - } - ], - } - controller.api.session_handler("data") + mock_unifi_websocket( + data={ + "meta": {"rc": "ok", "message": "events"}, + "data": [ + { + "datetime": "2020-01-20T19:37:04Z", + "key": aiounifi.events.WIRELESS_CLIENT_CONNECTED, + "msg": "User[11:22:33:44:55:66] has connected to WLAN", + "time": 1579549024893, + } + ], + }, + ) 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( diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index e8081a831c2..51dbd735e10 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/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 = { - "meta": {"message": MESSAGE_CLIENT}, - "data": [client_1_copy], - } - controller.api.session_handler(SIGNAL_DATA) + mock_unifi_websocket( + data={ + "meta": {"message": MESSAGE_CLIENT}, + "data": [client_1_copy], + } + ) 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}, - "data": [EVENT_CLIENT_1_WIRELESS_DISCONNECTED], - } - controller.api.session_handler(SIGNAL_DATA) + + mock_unifi_websocket( + data={ + "meta": {"message": MESSAGE_EVENT}, + "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 = { - "meta": {"message": MESSAGE_EVENT}, - "data": [EVENT_CLIENT_1_WIRELESS_CONNECTED], - } - controller.api.session_handler(SIGNAL_DATA) + mock_unifi_websocket( + data={ + "meta": {"message": MESSAGE_EVENT}, + "data": [EVENT_CLIENT_1_WIRELESS_CONNECTED], + } + ) 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]} - controller.api.message_handler(event) + mock_unifi_websocket( + 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]} - controller.api.message_handler(event) + mock_unifi_websocket( + 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]} - controller.api.message_handler(event) + mock_unifi_websocket( + 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]} - controller.api.message_handler(event) + mock_unifi_websocket( + 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]} - controller.api.message_handler(message) - event = {"meta": {"message": MESSAGE_EVENT}, "data": [EVENT_DEVICE_2_UPGRADED]} - controller.api.message_handler(event) + mock_unifi_websocket( + data={ + "meta": {"message": MESSAGE_DEVICE}, + "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 = { - "meta": {"message": MESSAGE_CLIENT_REMOVED}, - "data": [CLIENT_1], - } - controller.api.session_handler(SIGNAL_DATA) + mock_unifi_websocket( + data={ + "meta": {"message": MESSAGE_CLIENT_REMOVED}, + "data": [CLIENT_1], + } + ) 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( - SIGNAL_CONNECTION_STATE, STATE_DISCONNECTED - ) + mock_unifi_websocket(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]} - controller.api.message_handler(event) + mock_unifi_websocket( + 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]} - controller.api.message_handler(event) + mock_unifi_websocket( + 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) - event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_3_copy]} - controller.api.message_handler(event) + + mock_unifi_websocket( + data={ + "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]} - controller.api.message_handler(event) + mock_unifi_websocket( + 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]} - controller.api.message_handler(event) + mock_unifi_websocket( + 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]} - controller.api.message_handler(event) + mock_unifi_websocket( + 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]} - controller.api.message_handler(event) + mock_unifi_websocket( + 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]} - controller.api.message_handler(event) + mock_unifi_websocket( + 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]} - controller.api.message_handler(event) + mock_unifi_websocket( + 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]} - controller.api.message_handler(event) + mock_unifi_websocket( + 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]} - controller.api.message_handler(event) + mock_unifi_websocket( + data={ + "meta": {"message": MESSAGE_CLIENT}, + "data": [client_1_client], + } + ) await hass.async_block_till_done() # Client is wireless and still connected diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py index 6d8b894fc34..591165dabf2 100644 --- a/tests/components/unifi/test_init.py +++ b/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, - data=ENTRY_CONFIG, - unique_id="1", - version=1, + domain=UNIFI_DOMAIN, data=ENTRY_CONFIG, unique_id="1", entry_id=1 ) entry.add_to_hass(hass) - mock_registry = Mock() - with patch( - "homeassistant.components.unifi.UniFiController" - ) as mock_controller, patch( - "homeassistant.helpers.device_registry.async_get_registry", - return_value=mock_coro(mock_registry), - ): + + with patch("homeassistant.components.unifi.UniFiController") as mock_controller: 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" + ] + ) diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py index c668bf3789f..db1794e0878 100644 --- a/tests/components/unifi/test_sensor.py +++ b/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} - controller.api.message_handler(event) + mock_unifi_websocket( + 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 = { - "meta": {"message": MESSAGE_CLIENT_REMOVED}, - "data": [CLIENTS[0]], - } - controller.api.session_handler(SIGNAL_DATA) + mock_unifi_websocket( + data={ + "meta": {"message": MESSAGE_CLIENT_REMOVED}, + "data": [CLIENTS[0]], + } + ) await hass.async_block_till_done() assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 3 diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index e5a3a7eccc4..a0dc8e984e1 100644 --- a/tests/components/unifi/test_switch.py +++ b/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 = { - "meta": {"message": MESSAGE_CLIENT_REMOVED}, - "data": [CLIENT_1, UNBLOCKED], - } - controller.api.session_handler(SIGNAL_DATA) + mock_unifi_websocket( + data={ + "meta": {"message": MESSAGE_CLIENT_REMOVED}, + "data": [CLIENT_1, UNBLOCKED], + } + ) 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 = { - "meta": {"message": MESSAGE_EVENT}, - "data": [EVENT_BLOCKED_CLIENT_UNBLOCKED], - } - controller.api.session_handler(SIGNAL_DATA) + mock_unifi_websocket( + data={ + "meta": {"message": MESSAGE_EVENT}, + "data": [EVENT_BLOCKED_CLIENT_UNBLOCKED], + } + ) 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 = { - "meta": {"message": MESSAGE_EVENT}, - "data": [EVENT_BLOCKED_CLIENT_BLOCKED], - } - controller.api.session_handler(SIGNAL_DATA) + mock_unifi_websocket( + data={ + "meta": {"message": MESSAGE_EVENT}, + "data": [EVENT_BLOCKED_CLIENT_BLOCKED], + } + ) 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 = { - "meta": {"message": "sta:sync"}, - "data": [BLOCKED], - } - controller.api.session_handler(SIGNAL_DATA) + mock_unifi_websocket( + data={ + "meta": {"message": "sta:sync"}, + "data": [BLOCKED], + } + ) await hass.async_block_till_done() assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0 - controller.api.websocket._data = { - "meta": {"message": MESSAGE_EVENT}, - "data": [EVENT_BLOCKED_CLIENT_CONNECTED], - } - controller.api.session_handler(SIGNAL_DATA) + mock_unifi_websocket( + data={ + "meta": {"message": MESSAGE_EVENT}, + "data": [EVENT_BLOCKED_CLIENT_CONNECTED], + } + ) 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 = { - "meta": {"message": "sta:sync"}, - "data": [CLIENT_2], - } - controller.api.session_handler(SIGNAL_DATA) + mock_unifi_websocket( + data={ + "meta": {"message": "sta:sync"}, + "data": [CLIENT_2], + } + ) await hass.async_block_till_done() assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 - controller.api.websocket._data = { - "meta": {"message": MESSAGE_EVENT}, - "data": [EVENT_CLIENT_2_CONNECTED], - } - controller.api.session_handler(SIGNAL_DATA) + mock_unifi_websocket( + data={ + "meta": {"message": MESSAGE_EVENT}, + "data": [EVENT_CLIENT_2_CONNECTED], + } + ) await hass.async_block_till_done() assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2