"""Test UniFi Network.""" from copy import deepcopy from datetime import timedelta from http import HTTPStatus from unittest.mock import Mock, patch import aiounifi import pytest from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN from homeassistant.components.image import DOMAIN as IMAGE_DOMAIN from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN from homeassistant.components.unifi.const import ( CONF_SITE_ID, CONF_TRACK_CLIENTS, CONF_TRACK_DEVICES, DEFAULT_ALLOW_BANDWIDTH_SENSORS, DEFAULT_ALLOW_UPTIME_SENSORS, DEFAULT_DETECTION_TIME, DEFAULT_TRACK_CLIENTS, DEFAULT_TRACK_DEVICES, DEFAULT_TRACK_WIRED_CLIENTS, DOMAIN as UNIFI_DOMAIN, PLATFORMS, UNIFI_WIRELESS_CLIENTS, ) from homeassistant.components.unifi.errors import AuthenticationRequired, CannotConnect from homeassistant.components.unifi.hub import get_unifi_api from homeassistant.const import ( CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME, CONF_VERIFY_SSL, CONTENT_TYPE_JSON, ) from homeassistant.core import HomeAssistant from homeassistant.helpers import device_registry as dr 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.test_util.aiohttp import AiohttpClientMocker DEFAULT_CONFIG_ENTRY_ID = "1" DEFAULT_HOST = "1.2.3.4" DEFAULT_SITE = "site_id" CONTROLLER_HOST = { "hostname": "controller_host", "ip": DEFAULT_HOST, "is_wired": True, "last_seen": 1562600145, "mac": "10:00:00:00:00:01", "name": "Controller host", "oui": "Producer", "sw_mac": "00:00:00:00:01:01", "sw_port": 1, "wired-rx_bytes": 1234000000, "wired-tx_bytes": 5678000000, "uptime": 1562600160, } ENTRY_CONFIG = { CONF_HOST: DEFAULT_HOST, CONF_USERNAME: "username", CONF_PASSWORD: "password", CONF_PORT: 1234, CONF_SITE_ID: DEFAULT_SITE, CONF_VERIFY_SSL: False, } ENTRY_OPTIONS = {} CONFIGURATION = [] SITE = [{"desc": "Site name", "name": "site_id", "role": "admin", "_id": "1"}] SYSTEM_INFORMATION = [ { "anonymous_controller_id": "24f81231-a456-4c32-abcd-f5612345385f", "build": "atag_7.4.162_21057", "console_display_version": "3.1.15", "hostname": "UDMP", "name": "UDMP", "previous_version": "7.4.156", "timezone": "Europe/Stockholm", "ubnt_device_type": "UDMPRO", "udm_version": "3.0.20.9281", "update_available": False, "update_downloaded": False, "uptime": 1196290, "version": "7.4.162", } ] def mock_default_unifi_requests( aioclient_mock, host, site_id, sites=None, clients_response=None, clients_all_response=None, devices_response=None, dpiapp_response=None, dpigroup_response=None, port_forward_response=None, system_information_response=None, wlans_response=None, ): """Mock default UniFi requests responses.""" aioclient_mock.get(f"https://{host}:1234", status=302) # Check UniFi OS aioclient_mock.post( f"https://{host}:1234/api/login", json={"data": "login successful", "meta": {"rc": "ok"}}, headers={"content-type": CONTENT_TYPE_JSON}, ) aioclient_mock.get( f"https://{host}:1234/api/self/sites", json={"data": sites or [], "meta": {"rc": "ok"}}, headers={"content-type": CONTENT_TYPE_JSON}, ) aioclient_mock.get( f"https://{host}:1234/api/s/{site_id}/stat/sta", json={"data": clients_response or [], "meta": {"rc": "ok"}}, headers={"content-type": CONTENT_TYPE_JSON}, ) aioclient_mock.get( f"https://{host}:1234/api/s/{site_id}/rest/user", json={"data": clients_all_response or [], "meta": {"rc": "ok"}}, headers={"content-type": CONTENT_TYPE_JSON}, ) aioclient_mock.get( f"https://{host}:1234/api/s/{site_id}/stat/device", json={"data": devices_response or [], "meta": {"rc": "ok"}}, headers={"content-type": CONTENT_TYPE_JSON}, ) aioclient_mock.get( f"https://{host}:1234/api/s/{site_id}/rest/dpiapp", json={"data": dpiapp_response or [], "meta": {"rc": "ok"}}, headers={"content-type": CONTENT_TYPE_JSON}, ) aioclient_mock.get( f"https://{host}:1234/api/s/{site_id}/rest/dpigroup", json={"data": dpigroup_response or [], "meta": {"rc": "ok"}}, headers={"content-type": CONTENT_TYPE_JSON}, ) aioclient_mock.get( f"https://{host}:1234/api/s/{site_id}/rest/portforward", json={"data": port_forward_response or [], "meta": {"rc": "ok"}}, headers={"content-type": CONTENT_TYPE_JSON}, ) aioclient_mock.get( f"https://{host}:1234/api/s/{site_id}/stat/sysinfo", json={"data": system_information_response or [], "meta": {"rc": "ok"}}, headers={"content-type": CONTENT_TYPE_JSON}, ) aioclient_mock.get( f"https://{host}:1234/api/s/{site_id}/rest/wlanconf", json={"data": wlans_response or [], "meta": {"rc": "ok"}}, headers={"content-type": CONTENT_TYPE_JSON}, ) aioclient_mock.get( f"https://{host}:1234/v2/api/site/{site_id}/trafficroutes", json=[{}], headers={"content-type": CONTENT_TYPE_JSON}, ) aioclient_mock.get( f"https://{host}:1234/v2/api/site/{site_id}/trafficrules", json=[{}], headers={"content-type": CONTENT_TYPE_JSON}, ) async def setup_unifi_integration( hass, aioclient_mock=None, *, config=ENTRY_CONFIG, options=ENTRY_OPTIONS, sites=SITE, clients_response=None, clients_all_response=None, devices_response=None, dpiapp_response=None, dpigroup_response=None, port_forward_response=None, system_information_response=None, wlans_response=None, known_wireless_clients=None, unique_id="1", config_entry_id=DEFAULT_CONFIG_ENTRY_ID, ): """Create the UniFi Network instance.""" assert await async_setup_component(hass, UNIFI_DOMAIN, {}) config_entry = MockConfigEntry( domain=UNIFI_DOMAIN, data=deepcopy(config), options=deepcopy(options), unique_id=unique_id, entry_id=config_entry_id, version=1, ) config_entry.add_to_hass(hass) if known_wireless_clients: hass.data[UNIFI_WIRELESS_CLIENTS].wireless_clients.update( known_wireless_clients ) if aioclient_mock: mock_default_unifi_requests( aioclient_mock, host=config_entry.data[CONF_HOST], site_id=config_entry.data[CONF_SITE_ID], sites=sites, clients_response=clients_response, clients_all_response=clients_all_response, devices_response=devices_response, dpiapp_response=dpiapp_response, dpigroup_response=dpigroup_response, port_forward_response=port_forward_response, system_information_response=system_information_response, wlans_response=wlans_response, ) 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]: return None return config_entry async def test_hub_setup( hass: HomeAssistant, device_registry: dr.DeviceRegistry, aioclient_mock: AiohttpClientMocker, ) -> None: """Successful setup.""" with patch( "homeassistant.config_entries.ConfigEntries.async_forward_entry_setup", return_value=True, ) as forward_entry_setup: config_entry = await setup_unifi_integration( hass, aioclient_mock, system_information_response=SYSTEM_INFORMATION ) hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id] entry = hub.config.entry assert len(forward_entry_setup.mock_calls) == len(PLATFORMS) assert forward_entry_setup.mock_calls[0][1] == (entry, BUTTON_DOMAIN) assert forward_entry_setup.mock_calls[1][1] == (entry, TRACKER_DOMAIN) assert forward_entry_setup.mock_calls[2][1] == (entry, IMAGE_DOMAIN) assert forward_entry_setup.mock_calls[3][1] == (entry, SENSOR_DOMAIN) assert forward_entry_setup.mock_calls[4][1] == (entry, SWITCH_DOMAIN) assert hub.config.host == ENTRY_CONFIG[CONF_HOST] assert hub.is_admin == (SITE[0]["role"] == "admin") assert hub.config.option_allow_bandwidth_sensors == DEFAULT_ALLOW_BANDWIDTH_SENSORS assert hub.config.option_allow_uptime_sensors == DEFAULT_ALLOW_UPTIME_SENSORS assert isinstance(hub.config.option_block_clients, list) assert hub.config.option_track_clients == DEFAULT_TRACK_CLIENTS assert hub.config.option_track_devices == DEFAULT_TRACK_DEVICES assert hub.config.option_track_wired_clients == DEFAULT_TRACK_WIRED_CLIENTS assert hub.config.option_detection_time == timedelta(seconds=DEFAULT_DETECTION_TIME) assert isinstance(hub.config.option_ssid_filter, set) assert hub.signal_reachable == "unifi-reachable-1" assert hub.signal_options_update == "unifi-options-1" assert hub.signal_heartbeat_missed == "unifi-heartbeat-missed" device_entry = device_registry.async_get_or_create( config_entry_id=config_entry.entry_id, identifiers={(UNIFI_DOMAIN, config_entry.unique_id)}, ) assert device_entry.sw_version == "7.4.162" async def test_hub_not_accessible(hass: HomeAssistant) -> None: """Retry to login gets scheduled when connection fails.""" with patch( "homeassistant.components.unifi.hub.get_unifi_api", side_effect=CannotConnect, ): await setup_unifi_integration(hass) assert hass.data[UNIFI_DOMAIN] == {} async def test_hub_trigger_reauth_flow(hass: HomeAssistant) -> None: """Failed authentication trigger a reauthentication flow.""" with patch( "homeassistant.components.unifi.get_unifi_api", side_effect=AuthenticationRequired, ), patch.object(hass.config_entries.flow, "async_init") as mock_flow_init: await setup_unifi_integration(hass) mock_flow_init.assert_called_once() assert hass.data[UNIFI_DOMAIN] == {} async def test_hub_unknown_error(hass: HomeAssistant) -> None: """Unknown errors are handled.""" with patch( "homeassistant.components.unifi.hub.get_unifi_api", side_effect=Exception, ): await setup_unifi_integration(hass) assert hass.data[UNIFI_DOMAIN] == {} async def test_config_entry_updated( hass: HomeAssistant, aioclient_mock: AiohttpClientMocker ) -> None: """Calling reset when the entry has been setup.""" config_entry = await setup_unifi_integration(hass, aioclient_mock) hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id] event_call = Mock() unsub = async_dispatcher_connect(hass, hub.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: HomeAssistant, aioclient_mock: AiohttpClientMocker ) -> None: """Calling reset when the entry has been setup.""" config_entry = await setup_unifi_integration(hass, aioclient_mock) hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id] result = await hub.async_reset() await hass.async_block_till_done() assert result is True async def test_reset_fails( hass: HomeAssistant, aioclient_mock: AiohttpClientMocker ) -> None: """Calling reset when the entry has been setup can return false.""" config_entry = await setup_unifi_integration(hass, aioclient_mock) hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id] with patch( "homeassistant.config_entries.ConfigEntries.async_forward_entry_unload", return_value=False, ): result = await hub.async_reset() await hass.async_block_till_done() assert result is False async def test_connection_state_signalling( hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_device_registry, websocket_mock, ) -> None: """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" await websocket_mock.disconnect() # Controller is disconnected assert hass.states.get("device_tracker.client").state == "unavailable" await websocket_mock.reconnect() # Controller is once again connected assert hass.states.get("device_tracker.client").state == "home" async def test_reconnect_mechanism( hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock ) -> None: """Verify reconnect prints only on first reconnection try.""" await setup_unifi_integration(hass, aioclient_mock) aioclient_mock.clear_requests() aioclient_mock.get(f"https://{DEFAULT_HOST}:1234/", status=HTTPStatus.BAD_GATEWAY) await websocket_mock.disconnect() assert aioclient_mock.call_count == 0 await websocket_mock.reconnect(fail=True) assert aioclient_mock.call_count == 1 await websocket_mock.reconnect(fail=True) assert aioclient_mock.call_count == 2 @pytest.mark.parametrize( "exception", [ TimeoutError, aiounifi.BadGateway, aiounifi.ServiceUnavailable, aiounifi.AiounifiException, ], ) async def test_reconnect_mechanism_exceptions( hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock, exception ) -> None: """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.hub.hub.UnifiWebsocket.reconnect" ) as mock_reconnect: await websocket_mock.disconnect() await websocket_mock.reconnect() mock_reconnect.assert_called_once() async def test_get_unifi_api(hass: HomeAssistant) -> None: """Successful call.""" with patch("aiounifi.Controller.login", return_value=True): assert await get_unifi_api(hass, ENTRY_CONFIG) async def test_get_unifi_api_verify_ssl_false(hass: HomeAssistant) -> None: """Successful call with verify ssl set to false.""" hub_data = dict(ENTRY_CONFIG) hub_data[CONF_VERIFY_SSL] = False with patch("aiounifi.Controller.login", return_value=True): assert await get_unifi_api(hass, hub_data) @pytest.mark.parametrize( ("side_effect", "raised_exception"), [ (TimeoutError, CannotConnect), (aiounifi.BadGateway, CannotConnect), (aiounifi.Forbidden, CannotConnect), (aiounifi.ServiceUnavailable, CannotConnect), (aiounifi.RequestError, CannotConnect), (aiounifi.ResponseError, CannotConnect), (aiounifi.Unauthorized, AuthenticationRequired), (aiounifi.LoginRequired, AuthenticationRequired), (aiounifi.AiounifiException, AuthenticationRequired), ], ) async def test_get_unifi_api_fails_to_connect( hass: HomeAssistant, side_effect, raised_exception ) -> None: """Check that get_unifi_api can handle UniFi Network being unavailable.""" with patch("aiounifi.Controller.login", side_effect=side_effect), pytest.raises( raised_exception ): await get_unifi_api(hass, ENTRY_CONFIG)