diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index 3a2b37487af..8e05d8a1dd1 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -4,12 +4,7 @@ from copy import copy from datetime import timedelta -from asynctest import Mock - -import pytest - -from aiounifi.clients import Clients, ClientsAll -from aiounifi.devices import Devices +from asynctest import patch from homeassistant import config_entries from homeassistant.components import unifi @@ -107,64 +102,63 @@ ENTRY_CONFIG = {CONF_CONTROLLER: CONTROLLER_DATA} CONTROLLER_ID = CONF_CONTROLLER_ID.format(host="mock-host", site="mock-site") -@pytest.fixture -def mock_controller(hass): - """Mock a UniFi Controller.""" - hass.data[UNIFI_CONFIG] = {} - hass.data[UNIFI_WIRELESS_CLIENTS] = Mock() - controller = unifi.UniFiController(hass, None) - controller.wireless_clients = set() - - controller.api = Mock() - controller.mock_requests = [] - - controller.mock_client_responses = deque() - controller.mock_device_responses = deque() - controller.mock_client_all_responses = deque() - - async def mock_request(method, path, **kwargs): - kwargs["method"] = method - kwargs["path"] = path - controller.mock_requests.append(kwargs) - if path == "s/{site}/stat/sta": - return controller.mock_client_responses.popleft() - if path == "s/{site}/stat/device": - return controller.mock_device_responses.popleft() - if path == "s/{site}/rest/user": - return controller.mock_client_all_responses.popleft() - return None - - controller.api.clients = Clients({}, mock_request) - controller.api.devices = Devices({}, mock_request) - controller.api.clients_all = ClientsAll({}, mock_request) - - return controller - - -async def setup_controller(hass, mock_controller, options={}): - """Load the UniFi switch platform with the provided controller.""" - hass.config.components.add(unifi.DOMAIN) - hass.data[unifi.DOMAIN] = {CONTROLLER_ID: mock_controller} +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( - 1, - unifi.DOMAIN, - "Mock Title", - ENTRY_CONFIG, - "test", - config_entries.CONN_CLASS_LOCAL_POLL, - entry_id=1, + version=1, + domain=unifi.DOMAIN, + title="Mock Title", + data=config, + source="test", + connection_class=config_entries.CONN_CLASS_LOCAL_POLL, system_options={}, options=options, - ) - hass.config_entries._entries.append(config_entry) - mock_controller.config_entry = config_entry - - await mock_controller.async_update() - await hass.config_entries.async_forward_entry_setup( - config_entry, device_tracker.DOMAIN + 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}) + print(mock_requests, mock_client_responses, mock_device_responses) + 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): @@ -178,24 +172,30 @@ async def test_platform_manually_configured(hass): assert unifi.DOMAIN not in hass.data -async def test_no_clients(hass, mock_controller): +async def test_no_clients(hass): """Test the update_clients function when no clients are found.""" - mock_controller.mock_client_responses.append({}) - mock_controller.mock_device_responses.append({}) + await setup_unifi_integration( + hass, + ENTRY_CONFIG, + options={}, + clients_response={}, + devices_response={}, + clients_all_response={}, + ) - await setup_controller(hass, mock_controller) - assert len(mock_controller.mock_requests) == 2 assert len(hass.states.async_all()) == 2 -async def test_tracked_devices(hass, mock_controller): +async def test_tracked_devices(hass): """Test the update_items function with some clients.""" - mock_controller.mock_client_responses.append([CLIENT_1, CLIENT_2, CLIENT_3]) - mock_controller.mock_device_responses.append([DEVICE_1, DEVICE_2]) - options = {CONF_SSID_FILTER: ["ssid"]} - - await setup_controller(hass, mock_controller, options) - assert len(mock_controller.mock_requests) == 2 + controller = await setup_unifi_integration( + hass, + ENTRY_CONFIG, + options={CONF_SSID_FILTER: ["ssid"]}, + clients_response=[CLIENT_1, CLIENT_2, CLIENT_3], + devices_response=[DEVICE_1, DEVICE_2], + clients_all_response={}, + ) assert len(hass.states.async_all()) == 5 client_1 = hass.states.get("device_tracker.client_1") @@ -217,9 +217,9 @@ async def test_tracked_devices(hass, mock_controller): client_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow()) device_1_copy = copy(DEVICE_1) device_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow()) - mock_controller.mock_client_responses.append([client_1_copy]) - mock_controller.mock_device_responses.append([device_1_copy]) - await mock_controller.async_update() + controller.mock_client_responses.append([client_1_copy]) + controller.mock_device_responses.append([device_1_copy]) + await controller.async_update() await hass.async_block_till_done() client_1 = hass.states.get("device_tracker.client_1") @@ -230,19 +230,17 @@ async def test_tracked_devices(hass, mock_controller): device_1_copy = copy(DEVICE_1) device_1_copy["disabled"] = True - mock_controller.mock_client_responses.append({}) - mock_controller.mock_device_responses.append([device_1_copy]) - await mock_controller.async_update() + controller.mock_client_responses.append({}) + controller.mock_device_responses.append([device_1_copy]) + await controller.async_update() await hass.async_block_till_done() device_1 = hass.states.get("device_tracker.device_1") assert device_1.state == STATE_UNAVAILABLE - mock_controller.config_entry.add_update_listener( - mock_controller.async_options_updated - ) + controller.config_entry.add_update_listener(controller.async_options_updated) hass.config_entries.async_update_entry( - mock_controller.config_entry, + controller.config_entry, options={ CONF_SSID_FILTER: [], CONF_TRACK_WIRED_CLIENTS: False, @@ -258,18 +256,22 @@ async def test_tracked_devices(hass, mock_controller): assert device_1 is None -async def test_wireless_client_go_wired_issue(hass, mock_controller): +async def test_wireless_client_go_wired_issue(hass): """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. """ client_1_client = copy(CLIENT_1) client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow()) - mock_controller.mock_client_responses.append([client_1_client]) - mock_controller.mock_device_responses.append({}) - await setup_controller(hass, mock_controller) - assert len(mock_controller.mock_requests) == 2 + controller = await setup_unifi_integration( + hass, + ENTRY_CONFIG, + options={}, + clients_response=[client_1_client], + devices_response={}, + clients_all_response={}, + ) assert len(hass.states.async_all()) == 3 client_1 = hass.states.get("device_tracker.client_1") @@ -278,9 +280,9 @@ async def test_wireless_client_go_wired_issue(hass, mock_controller): client_1_client["is_wired"] = True client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow()) - mock_controller.mock_client_responses.append([client_1_client]) - mock_controller.mock_device_responses.append({}) - await mock_controller.async_update() + controller.mock_client_responses.append([client_1_client]) + controller.mock_device_responses.append({}) + await controller.async_update() await hass.async_block_till_done() client_1 = hass.states.get("device_tracker.client_1") @@ -288,31 +290,27 @@ async def test_wireless_client_go_wired_issue(hass, mock_controller): client_1_client["is_wired"] = False client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow()) - mock_controller.mock_client_responses.append([client_1_client]) - mock_controller.mock_device_responses.append({}) - await mock_controller.async_update() + controller.mock_client_responses.append([client_1_client]) + controller.mock_device_responses.append({}) + await controller.async_update() await hass.async_block_till_done() client_1 = hass.states.get("device_tracker.client_1") assert client_1.state == "home" -async def test_restoring_client(hass, mock_controller): +async def test_restoring_client(hass): """Test the update_items function with some clients.""" - mock_controller.mock_client_responses.append([CLIENT_2]) - mock_controller.mock_device_responses.append({}) - mock_controller.mock_client_all_responses.append([CLIENT_1]) - options = {unifi.CONF_BLOCK_CLIENT: True} - config_entry = config_entries.ConfigEntry( - 1, - unifi.DOMAIN, - "Mock Title", - ENTRY_CONFIG, - "test", - config_entries.CONN_CLASS_LOCAL_POLL, - entry_id=1, + version=1, + domain=unifi.DOMAIN, + title="Mock Title", + data=ENTRY_CONFIG, + source="test", + connection_class=config_entries.CONN_CLASS_LOCAL_POLL, system_options={}, + options={}, + entry_id=1, ) registry = await entity_registry.async_get_registry(hass) @@ -331,22 +329,30 @@ async def test_restoring_client(hass, mock_controller): config_entry=config_entry, ) - await setup_controller(hass, mock_controller, options) - assert len(mock_controller.mock_requests) == 3 + await setup_unifi_integration( + hass, + ENTRY_CONFIG, + options={unifi.CONF_BLOCK_CLIENT: True}, + clients_response=[CLIENT_2], + devices_response={}, + clients_all_response=[CLIENT_1], + ) assert len(hass.states.async_all()) == 4 device_1 = hass.states.get("device_tracker.client_1") assert device_1 is not None -async def test_dont_track_clients(hass, mock_controller): +async def test_dont_track_clients(hass): """Test dont track clients config works.""" - mock_controller.mock_client_responses.append([CLIENT_1]) - mock_controller.mock_device_responses.append([DEVICE_1]) - options = {unifi.controller.CONF_TRACK_CLIENTS: False} - - await setup_controller(hass, mock_controller, options) - assert len(mock_controller.mock_requests) == 2 + await setup_unifi_integration( + hass, + ENTRY_CONFIG, + options={unifi.controller.CONF_TRACK_CLIENTS: False}, + clients_response=[CLIENT_1], + devices_response=[DEVICE_1], + clients_all_response={}, + ) assert len(hass.states.async_all()) == 3 client_1 = hass.states.get("device_tracker.client_1") @@ -357,14 +363,16 @@ async def test_dont_track_clients(hass, mock_controller): assert device_1.state == "not_home" -async def test_dont_track_devices(hass, mock_controller): +async def test_dont_track_devices(hass): """Test dont track devices config works.""" - mock_controller.mock_client_responses.append([CLIENT_1]) - mock_controller.mock_device_responses.append([DEVICE_1]) - options = {unifi.controller.CONF_TRACK_DEVICES: False} - - await setup_controller(hass, mock_controller, options) - assert len(mock_controller.mock_requests) == 2 + await setup_unifi_integration( + hass, + ENTRY_CONFIG, + options={unifi.controller.CONF_TRACK_DEVICES: False}, + clients_response=[CLIENT_1], + devices_response=[DEVICE_1], + clients_all_response={}, + ) assert len(hass.states.async_all()) == 3 client_1 = hass.states.get("device_tracker.client_1") @@ -375,14 +383,16 @@ async def test_dont_track_devices(hass, mock_controller): assert device_1 is None -async def test_dont_track_wired_clients(hass, mock_controller): +async def test_dont_track_wired_clients(hass): """Test dont track wired clients config works.""" - mock_controller.mock_client_responses.append([CLIENT_1, CLIENT_2]) - mock_controller.mock_device_responses.append({}) - options = {unifi.controller.CONF_TRACK_WIRED_CLIENTS: False} - - await setup_controller(hass, mock_controller, options) - assert len(mock_controller.mock_requests) == 2 + await setup_unifi_integration( + hass, + ENTRY_CONFIG, + options={unifi.controller.CONF_TRACK_WIRED_CLIENTS: False}, + clients_response=[CLIENT_1, CLIENT_2], + devices_response={}, + clients_all_response={}, + ) assert len(hass.states.async_all()) == 3 client_1 = hass.states.get("device_tracker.client_1")