"""Fixtures for UniFi Network methods.""" from __future__ import annotations import asyncio from collections.abc import Callable from datetime import timedelta from types import MappingProxyType from typing import Any from unittest.mock import patch from aiounifi.models.message import MessageKey import pytest from homeassistant.components.unifi.const import CONF_SITE_ID, DOMAIN as UNIFI_DOMAIN from homeassistant.components.unifi.hub.websocket import RETRY_TIMER from homeassistant.config_entries import ConfigEntry 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 import homeassistant.util.dt as dt_util from tests.common import MockConfigEntry, async_fire_time_changed from tests.test_util.aiohttp import AiohttpClientMocker DEFAULT_CONFIG_ENTRY_ID = "1" DEFAULT_HOST = "1.2.3.4" DEFAULT_PORT = 1234 DEFAULT_SITE = "site_id" @pytest.fixture(autouse=True) def mock_discovery(): """No real network traffic allowed.""" with patch( "homeassistant.components.unifi.config_flow._async_discover_unifi", return_value=None, ) as mock: yield mock @pytest.fixture def mock_device_registry(hass, device_registry: dr.DeviceRegistry): """Mock device registry.""" config_entry = MockConfigEntry(domain="something_else") config_entry.add_to_hass(hass) for idx, device in enumerate( ( "00:00:00:00:00:01", "00:00:00:00:00:02", "00:00:00:00:00:03", "00:00:00:00:00:04", "00:00:00:00:00:05", "00:00:00:00:00:06", "00:00:00:00:01:01", "00:00:00:00:02:02", ) ): device_registry.async_get_or_create( name=f"Device {idx}", config_entry_id=config_entry.entry_id, connections={(dr.CONNECTION_NETWORK_MAC, device)}, ) # Config entry fixtures @pytest.fixture(name="config_entry") def config_entry_fixture( hass: HomeAssistant, config_entry_data: MappingProxyType[str, Any], config_entry_options: MappingProxyType[str, Any], ) -> ConfigEntry: """Define a config entry fixture.""" config_entry = MockConfigEntry( domain=UNIFI_DOMAIN, entry_id="1", unique_id="1", data=config_entry_data, options=config_entry_options, ) config_entry.add_to_hass(hass) return config_entry @pytest.fixture(name="config_entry_data") def config_entry_data_fixture() -> MappingProxyType[str, Any]: """Define a config entry data fixture.""" return { CONF_HOST: DEFAULT_HOST, CONF_USERNAME: "username", CONF_PASSWORD: "password", CONF_PORT: DEFAULT_PORT, CONF_SITE_ID: DEFAULT_SITE, CONF_VERIFY_SSL: False, } @pytest.fixture(name="config_entry_options") def config_entry_options_fixture() -> MappingProxyType[str, Any]: """Define a config entry options fixture.""" return {} @pytest.fixture(name="mock_requests") def request_fixture( aioclient_mock: AiohttpClientMocker, client_payload: list[dict[str, Any]], clients_all_payload: list[dict[str, Any]], device_payload: list[dict[str, Any]], dpi_app_payload: list[dict[str, Any]], dpi_group_payload: list[dict[str, Any]], port_forward_payload: list[dict[str, Any]], site_payload: list[dict[str, Any]], system_information_payload: list[dict[str, Any]], wlan_payload: list[dict[str, Any]], ) -> Callable[[str], None]: """Mock default UniFi requests responses.""" def __mock_requests(host: str = DEFAULT_HOST, site_id: str = DEFAULT_SITE) -> None: url = f"https://{host}:{DEFAULT_PORT}" def mock_get_request(path: str, payload: list[dict[str, Any]]) -> None: aioclient_mock.get( f"{url}{path}", json={"meta": {"rc": "OK"}, "data": payload}, headers={"content-type": CONTENT_TYPE_JSON}, ) aioclient_mock.get(url, status=302) # UniFI OS check aioclient_mock.post( f"{url}/api/login", json={"data": "login successful", "meta": {"rc": "ok"}}, headers={"content-type": CONTENT_TYPE_JSON}, ) mock_get_request("/api/self/sites", site_payload) mock_get_request(f"/api/s/{site_id}/stat/sta", client_payload) mock_get_request(f"/api/s/{site_id}/rest/user", clients_all_payload) mock_get_request(f"/api/s/{site_id}/stat/device", device_payload) mock_get_request(f"/api/s/{site_id}/rest/dpiapp", dpi_app_payload) mock_get_request(f"/api/s/{site_id}/rest/dpigroup", dpi_group_payload) mock_get_request(f"/api/s/{site_id}/rest/portforward", port_forward_payload) mock_get_request(f"/api/s/{site_id}/stat/sysinfo", system_information_payload) mock_get_request(f"/api/s/{site_id}/rest/wlanconf", wlan_payload) return __mock_requests # Request payload fixtures @pytest.fixture(name="client_payload") def client_data_fixture() -> list[dict[str, Any]]: """Client data.""" return [] @pytest.fixture(name="clients_all_payload") def clients_all_data_fixture() -> list[dict[str, Any]]: """Clients all data.""" return [] @pytest.fixture(name="device_payload") def device_data_fixture() -> list[dict[str, Any]]: """Device data.""" return [] @pytest.fixture(name="dpi_app_payload") def dpi_app_data_fixture() -> list[dict[str, Any]]: """DPI app data.""" return [] @pytest.fixture(name="dpi_group_payload") def dpi_group_data_fixture() -> list[dict[str, Any]]: """DPI group data.""" return [] @pytest.fixture(name="port_forward_payload") def port_forward_data_fixture() -> list[dict[str, Any]]: """Port forward data.""" return [] @pytest.fixture(name="site_payload") def site_data_fixture() -> list[dict[str, Any]]: """Site data.""" return [{"desc": "Site name", "name": "site_id", "role": "admin", "_id": "1"}] @pytest.fixture(name="system_information_payload") def system_information_data_fixture() -> list[dict[str, Any]]: """System information data.""" return [ { "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", } ] @pytest.fixture(name="wlan_payload") def wlan_data_fixture() -> list[dict[str, Any]]: """WLAN data.""" return [] @pytest.fixture(name="mock_default_requests") def default_requests_fixture( mock_requests: Callable[[str, str], None], ) -> None: """Mock UniFi requests responses with default host and site.""" mock_requests(DEFAULT_HOST, DEFAULT_SITE) @pytest.fixture(name="config_entry_factory") async def config_entry_factory_fixture( hass: HomeAssistant, config_entry: ConfigEntry, mock_requests: Callable[[str, str], None], ) -> Callable[[], ConfigEntry]: """Fixture factory that can set up UniFi network integration.""" async def __mock_setup_config_entry() -> ConfigEntry: mock_requests(config_entry.data[CONF_HOST], config_entry.data[CONF_SITE_ID]) await hass.config_entries.async_setup(config_entry.entry_id) await hass.async_block_till_done() return config_entry return __mock_setup_config_entry @pytest.fixture(name="config_entry_setup") async def config_entry_setup_fixture( hass: HomeAssistant, config_entry_factory: Callable[[], ConfigEntry] ) -> ConfigEntry: """Fixture providing a set up instance of UniFi network integration.""" return await config_entry_factory() # Websocket fixtures class WebsocketStateManager(asyncio.Event): """Keep an async event that simules websocket context manager. Prepares disconnect and reconnect flows. """ def __init__(self, hass: HomeAssistant, aioclient_mock: AiohttpClientMocker): """Store hass object and initialize asyncio.Event.""" self.hass = hass self.aioclient_mock = aioclient_mock super().__init__() async def disconnect(self): """Mark future as done to make 'await self.api.start_websocket' return.""" self.set() await self.hass.async_block_till_done() async def reconnect(self, fail=False): """Set up new future to make 'await self.api.start_websocket' block. Mock api calls done by 'await self.api.login'. Fail will make 'await self.api.start_websocket' return immediately. """ hub = self.hass.config_entries.async_get_entry( DEFAULT_CONFIG_ENTRY_ID ).runtime_data self.aioclient_mock.get( f"https://{hub.config.host}:1234", status=302 ) # Check UniFi OS self.aioclient_mock.post( f"https://{hub.config.host}:1234/api/login", json={"data": "login successful", "meta": {"rc": "ok"}}, headers={"content-type": CONTENT_TYPE_JSON}, ) if not fail: self.clear() new_time = dt_util.utcnow() + timedelta(seconds=RETRY_TIMER) async_fire_time_changed(self.hass, new_time) await self.hass.async_block_till_done() @pytest.fixture(autouse=True) def websocket_mock(hass: HomeAssistant, aioclient_mock: AiohttpClientMocker): """Mock 'await self.api.start_websocket' in 'UniFiController.start_websocket'.""" websocket_state_manager = WebsocketStateManager(hass, aioclient_mock) with patch("aiounifi.Controller.start_websocket") as ws_mock: ws_mock.side_effect = websocket_state_manager.wait yield websocket_state_manager @pytest.fixture(autouse=True) def mock_unifi_websocket(hass): """No real websocket allowed.""" def make_websocket_call( *, message: MessageKey | None = None, data: list[dict] | dict | None = None, ): """Generate a websocket call.""" hub = hass.config_entries.async_get_entry(DEFAULT_CONFIG_ENTRY_ID).runtime_data if data and not message: hub.api.messages.handler(data) elif data and message: if not isinstance(data, list): data = [data] hub.api.messages.handler( { "meta": {"message": message.value}, "data": data, } ) else: raise NotImplementedError return make_websocket_call