core/tests/components/unifi/conftest.py

"""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,
        version=1,
    )
    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_unifi_requests")
def default_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_default_requests(host: str, site_id: str) -> 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_default_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="setup_default_unifi_requests")
def default_vapix_requests_fixture(
    config_entry: ConfigEntry,
    mock_unifi_requests: Callable[[str, str], None],
) -> None:
    """Mock default UniFi requests responses."""
    mock_unifi_requests(config_entry.data[CONF_HOST], config_entry.data[CONF_SITE_ID])


@pytest.fixture(name="prepare_config_entry")
async def prep_config_entry_fixture(
    hass: HomeAssistant, config_entry: ConfigEntry, setup_default_unifi_requests: None
) -> Callable[[], ConfigEntry]:
    """Fixture factory to set up UniFi network integration."""

    async def __mock_setup_config_entry() -> ConfigEntry:
        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="setup_config_entry")
async def setup_config_entry_fixture(
    hass: HomeAssistant, prepare_config_entry: Callable[[], ConfigEntry]
) -> ConfigEntry:
    """Fixture to set up UniFi network integration."""
    return await prepare_config_entry()


# 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