"""Fixtures for UniFi Network methods.""" from __future__ import annotations from unittest.mock import patch from aiounifi.websocket import SIGNAL_CONNECTION_STATE, SIGNAL_DATA 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(data: dict | None = None, state: str = ""): """Generate a websocket call.""" if data: mock.return_value.data = data mock.call_args[1]["callback"](SIGNAL_DATA) elif state: mock.return_value.state = state mock.call_args[1]["callback"](SIGNAL_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)}, )