"""Fixtures for UniFi Network methods.""" from __future__ import annotations from unittest.mock import patch from aiounifi.models.message import MessageKey from aiounifi.websocket import WebsocketSignal, WebsocketState import pytest from homeassistant.helpers import device_registry as dr from tests.common import MockConfigEntry @pytest.fixture(autouse=True) def mock_unifi_websocket(): """No real websocket allowed.""" with patch("aiounifi.controller.WSClient") as mock: def make_websocket_call( *, message: MessageKey | None = None, data: list[dict] | dict | None = None, state: WebsocketState | None = None, ): """Generate a websocket call.""" if data and not message: mock.return_value.data = data mock.call_args[1]["callback"](WebsocketSignal.DATA) elif data and message: if not isinstance(data, list): data = [data] mock.return_value.data = { "meta": {"message": message.value}, "data": data, } mock.call_args[1]["callback"](WebsocketSignal.DATA) elif state: mock.return_value.state = state mock.call_args[1]["callback"](WebsocketSignal.CONNECTION_STATE) else: raise NotImplementedError yield make_websocket_call @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): """Mock device registry.""" dev_reg = dr.async_get(hass) config_entry = MockConfigEntry(domain="something_else") 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:01:01", "00:00:00:00:02:02", ) ): dev_reg.async_get_or_create( name=f"Device {idx}", config_entry_id=config_entry.entry_id, connections={(dr.CONNECTION_NETWORK_MAC, device)}, )