jeans/core
jeans/core
1
0
Fork 0
mirror of https://github.com/home-assistant/core.git synced 2025-04-25 09:47:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Rework UniFi websocket (#100614)

Browse Source 
* Rework websocket management

* remove unnecessary fixture

* Remove controller from mock_unifi_websocket

* Mock api.login in reconnect method

* Remove unnecessary edits

* Minor clean up

* Bump aiounifi to v63

* Wait on task cancellation
This commit is contained in:
Robert Svensson 2023-09-27 10:56:24 +02:00 committed by GitHub
parent 134c005168
commit 01b5854968
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
13 changed files with 184 additions and 157 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
homeassistant/components/unifi/__init__.py
View File

@ -52,7 +52,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
if len(hass.data[UNIFI_DOMAIN]) == 1:
async_setup_services(hass)
api.start_websocket()
controller.start_websocket()
config_entry.async_on_unload(
hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, controller.shutdown)

55
homeassistant/components/unifi/controller.py
View File

@ -12,7 +12,6 @@ import aiounifi
from aiounifi.interfaces.api_handlers import ItemEvent
from aiounifi.models.configuration import Configuration
from aiounifi.models.device import DeviceSetPoePortModeRequest
from aiounifi.websocket import WebsocketState
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import (
@ -81,7 +80,7 @@ class UniFiController:
self.config_entry = config_entry
self.api = api
api.ws_state_callback = self.async_unifi_ws_state_callback
self.ws_task: asyncio.Task | None = None
self.available = True
self.wireless_clients = hass.data[UNIFI_WIRELESS_CLIENTS]
@ -223,23 +222,6 @@ class UniFiController:
for description in descriptions:
async_load_entities(description)
@callback
def async_unifi_ws_state_callback(self, state: WebsocketState) -> None:
"""Handle messages back from UniFi library."""
if state == WebsocketState.DISCONNECTED and self.available:
LOGGER.warning("Lost connection to UniFi Network")
if (state == WebsocketState.RUNNING and not self.available) or (
state == WebsocketState.DISCONNECTED and self.available
):
self.available = state == WebsocketState.RUNNING
async_dispatcher_send(self.hass, self.signal_reachable)
if not self.available:
self.hass.loop.call_later(RETRY_TIMER, self.reconnect, True)
else:
LOGGER.info("Connected to UniFi Network")
@property
def signal_reachable(self) -> str:
"""Integration specific event to signal a change in connection status."""
@ -367,6 +349,19 @@ class UniFiController:
controller.load_config_entry_options()
async_dispatcher_send(hass, controller.signal_options_update)
@callback
def start_websocket(self) -> None:
"""Start up connection to websocket."""
async def _websocket_runner() -> None:
"""Start websocket."""
await self.api.start_websocket()
self.available = False
async_dispatcher_send(self.hass, self.signal_reachable)
self.hass.loop.call_later(RETRY_TIMER, self.reconnect, True)
self.ws_task = self.hass.loop.create_task(_websocket_runner())
@callback
def reconnect(self, log: bool = False) -> None:
"""Prepare to reconnect UniFi session."""
@ -379,7 +374,11 @@ class UniFiController:
try:
async with asyncio.timeout(5):
await self.api.login()
self.api.start_websocket()
self.start_websocket()
if not self.available:
self.available = True
async_dispatcher_send(self.hass, self.signal_reachable)
except (
asyncio.TimeoutError,
@ -395,7 +394,8 @@ class UniFiController:
Used as an argument to EventBus.async_listen_once.
"""
self.api.stop_websocket()
if self.ws_task is not None:
self.ws_task.cancel()
async def async_reset(self) -> bool:
"""Reset this controller to default state.
@ -403,7 +403,18 @@ class UniFiController:
Will cancel any scheduled setup retry and will unload
the config entry.
"""
self.api.stop_websocket()
if self.ws_task is not None:
self.ws_task.cancel()
_, pending = await asyncio.wait([self.ws_task], timeout=10)
if pending:
LOGGER.warning(
"Unloading %s (%s) config entry. Task %s did not complete in time",
self.config_entry.title,
self.config_entry.domain,
self.ws_task,
)
unload_ok = await self.hass.config_entries.async_unload_platforms(
self.config_entry, PLATFORMS

2
homeassistant/components/unifi/manifest.json
View File

@ -8,7 +8,7 @@
"iot_class": "local_push",
"loggers": ["aiounifi"],
"quality_scale": "platinum",
"requirements": ["aiounifi==62"],
"requirements": ["aiounifi==63"],
"ssdp": [
{
"manufacturer": "Ubiquiti Networks",

2
requirements_all.txt
View File

@ -363,7 +363,7 @@ aiosyncthing==0.5.1
aiotractive==0.5.6
# homeassistant.components.unifi
aiounifi==62
aiounifi==63
# homeassistant.components.vlc_telnet
aiovlc==0.1.0

2
requirements_test_all.txt
View File

@ -338,7 +338,7 @@ aiosyncthing==0.5.1
aiotractive==0.5.6
# homeassistant.components.unifi
aiounifi==62
aiounifi==63
# homeassistant.components.vlc_telnet
aiovlc==0.1.0

105
tests/components/unifi/conftest.py
View File

@ -1,47 +1,100 @@
"""Fixtures for UniFi Network methods."""
from __future__ import annotations
import asyncio
from datetime import timedelta
from unittest.mock import patch
from aiounifi.models.message import MessageKey
from aiounifi.websocket import WebsocketSignal, WebsocketState
import pytest
from homeassistant.components.unifi.const import DOMAIN as UNIFI_DOMAIN
from homeassistant.components.unifi.controller import RETRY_TIMER
from homeassistant.const import 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
from tests.common import MockConfigEntry, async_fire_time_changed
from tests.components.unifi.test_controller import DEFAULT_CONFIG_ENTRY_ID
from tests.test_util.aiohttp import AiohttpClientMocker
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.
"""
controller = self.hass.data[UNIFI_DOMAIN][DEFAULT_CONFIG_ENTRY_ID]
self.aioclient_mock.get(
f"https://{controller.host}:1234", status=302
) # Check UniFi OS
self.aioclient_mock.post(
f"https://{controller.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 mock_unifi_websocket():
"""No real websocket allowed."""
with patch("aiounifi.controller.WSClient") as mock:
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
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 = {
@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."""
controller = hass.data[UNIFI_DOMAIN][DEFAULT_CONFIG_ENTRY_ID]
if data and not message:
controller.api.messages.handler(data)
elif data and message:
if not isinstance(data, list):
data = [data]
controller.api.messages.handler(
{
"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
)
else:
raise NotImplementedError
yield make_websocket_call
return make_websocket_call
@pytest.fixture(autouse=True)

10
tests/components/unifi/test_button.py
View File

@ -1,7 +1,5 @@
"""UniFi Network button platform tests."""
from aiounifi.websocket import WebsocketState
from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, ButtonDeviceClass
from homeassistant.components.unifi.const import DOMAIN as UNIFI_DOMAIN
from homeassistant.const import ATTR_DEVICE_CLASS, STATE_UNAVAILABLE, EntityCategory
@ -14,7 +12,7 @@ from tests.test_util.aiohttp import AiohttpClientMocker
async def test_restart_device_button(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock
) -> None:
"""Test restarting device button."""
config_entry = await setup_unifi_integration(
@ -71,11 +69,9 @@ async def test_restart_device_button(
# Availability signalling
# Controller disconnects
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
await websocket_mock.disconnect()
assert hass.states.get("button.switch_restart").state == STATE_UNAVAILABLE
# Controller reconnects
mock_unifi_websocket(state=WebsocketState.RUNNING)
await hass.async_block_till_done()
await websocket_mock.reconnect()
assert hass.states.get("button.switch_restart").state != STATE_UNAVAILABLE

42
tests/components/unifi/test_controller.py
View File

@ -6,7 +6,6 @@ from http import HTTPStatus
from unittest.mock import Mock, patch
import aiounifi
from aiounifi.websocket import WebsocketState
import pytest
from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN
@ -28,7 +27,7 @@ from homeassistant.components.unifi.const import (
PLATFORMS,
UNIFI_WIRELESS_CLIENTS,
)
from homeassistant.components.unifi.controller import RETRY_TIMER, get_unifi_controller
from homeassistant.components.unifi.controller import get_unifi_controller
from homeassistant.components.unifi.errors import AuthenticationRequired, CannotConnect
from homeassistant.const import (
CONF_HOST,
@ -44,7 +43,7 @@ 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, async_fire_time_changed
from tests.common import MockConfigEntry
from tests.test_util.aiohttp import AiohttpClientMocker
DEFAULT_CONFIG_ENTRY_ID = "1"
@ -365,8 +364,8 @@ async def test_reset_fails(
async def test_connection_state_signalling(
hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
mock_device_registry,
websocket_mock,
) -> None:
"""Verify connection statesignalling and connection state are working."""
client = {
@ -381,21 +380,17 @@ async def test_connection_state_signalling(
# Controller is connected
assert hass.states.get("device_tracker.client").state == "home"
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
await websocket_mock.disconnect()
# Controller is disconnected
assert hass.states.get("device_tracker.client").state == "unavailable"
mock_unifi_websocket(state=WebsocketState.RUNNING)
await hass.async_block_till_done()
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, mock_unifi_websocket
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock
) -> None:
"""Verify reconnect prints only on first reconnection try."""
await setup_unifi_integration(hass, aioclient_mock)
@ -403,21 +398,13 @@ async def test_reconnect_mechanism(
aioclient_mock.clear_requests()
aioclient_mock.get(f"https://{DEFAULT_HOST}:1234/", status=HTTPStatus.BAD_GATEWAY)
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
await websocket_mock.disconnect()
assert aioclient_mock.call_count == 0
new_time = dt_util.utcnow() + timedelta(seconds=RETRY_TIMER)
async_fire_time_changed(hass, new_time)
await hass.async_block_till_done()
await websocket_mock.reconnect(fail=True)
assert aioclient_mock.call_count == 1
new_time = dt_util.utcnow() + timedelta(seconds=RETRY_TIMER)
async_fire_time_changed(hass, new_time)
await hass.async_block_till_done()
await websocket_mock.reconnect(fail=True)
assert aioclient_mock.call_count == 2
@ -431,10 +418,7 @@ async def test_reconnect_mechanism(
],
)
async def test_reconnect_mechanism_exceptions(
hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
exception,
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock, exception
) -> None:
"""Verify async_reconnect calls expected methods."""
await setup_unifi_integration(hass, aioclient_mock)
@ -442,11 +426,9 @@ async def test_reconnect_mechanism_exceptions(
with patch("aiounifi.Controller.login", side_effect=exception), patch(
"homeassistant.components.unifi.controller.UniFiController.reconnect"
) as mock_reconnect:
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
await websocket_mock.disconnect()
new_time = dt_util.utcnow() + timedelta(seconds=RETRY_TIMER)
async_fire_time_changed(hass, new_time)
await websocket_mock.reconnect()
mock_reconnect.assert_called_once()

13
tests/components/unifi/test_device_tracker.py
View File

@ -3,7 +3,6 @@ from datetime import timedelta
from unittest.mock import patch
from aiounifi.models.message import MessageKey
from aiounifi.websocket import WebsocketState
from freezegun.api import FrozenDateTimeFactory
from homeassistant import config_entries
@ -40,8 +39,8 @@ async def test_no_entities(
async def test_tracked_wireless_clients(
hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
mock_device_registry,
mock_unifi_websocket,
) -> None:
"""Verify tracking of wireless clients."""
client = {
@ -402,7 +401,7 @@ async def test_remove_clients(
async def test_controller_state_change(
hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
websocket_mock,
mock_device_registry,
) -> None:
"""Verify entities state reflect on controller becoming unavailable."""
@ -443,16 +442,12 @@ async def test_controller_state_change(
assert hass.states.get("device_tracker.device").state == STATE_HOME
# Controller unavailable
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
await websocket_mock.disconnect()
assert hass.states.get("device_tracker.client").state == STATE_UNAVAILABLE
assert hass.states.get("device_tracker.device").state == STATE_UNAVAILABLE
# Controller available
mock_unifi_websocket(state=WebsocketState.RUNNING)
await hass.async_block_till_done()
await websocket_mock.reconnect()
assert hass.states.get("device_tracker.client").state == STATE_NOT_HOME
assert hass.states.get("device_tracker.device").state == STATE_HOME

8
tests/components/unifi/test_image.py
View File

@ -5,7 +5,6 @@ from datetime import timedelta
from http import HTTPStatus
from aiounifi.models.message import MessageKey
from aiounifi.websocket import WebsocketState
from syrupy.assertion import SnapshotAssertion
from homeassistant.components.image import DOMAIN as IMAGE_DOMAIN
@ -65,6 +64,7 @@ async def test_wlan_qr_code(
hass_client: ClientSessionGenerator,
snapshot: SnapshotAssertion,
mock_unifi_websocket,
websocket_mock,
) -> None:
"""Test the update_clients function when no clients are found."""
await setup_unifi_integration(hass, aioclient_mock, wlans_response=[WLAN])
@ -121,13 +121,11 @@ async def test_wlan_qr_code(
# Availability signalling
# Controller disconnects
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
await websocket_mock.disconnect()
assert hass.states.get("image.ssid_1_qr_code").state == STATE_UNAVAILABLE
# Controller reconnects
mock_unifi_websocket(state=WebsocketState.RUNNING)
await hass.async_block_till_done()
await websocket_mock.reconnect()
assert hass.states.get("image.ssid_1_qr_code").state != STATE_UNAVAILABLE
# WLAN gets disabled

31
tests/components/unifi/test_sensor.py
View File

@ -4,7 +4,6 @@ from datetime import datetime, timedelta
from unittest.mock import patch
from aiounifi.models.message import MessageKey
from aiounifi.websocket import WebsocketState
import pytest
from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
@ -562,7 +561,10 @@ async def test_remove_sensors(
async def test_poe_port_switches(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
websocket_mock,
) -> None:
"""Test the update_items function with some clients."""
await setup_unifi_integration(hass, aioclient_mock, devices_response=[DEVICE_1])
@ -607,16 +609,16 @@ async def test_poe_port_switches(
# Availability signalling
# Controller disconnects
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
await websocket_mock.disconnect()
assert (
hass.states.get("sensor.mock_name_port_1_poe_power").state == STATE_UNAVAILABLE
)
# Controller reconnects
mock_unifi_websocket(state=WebsocketState.RUNNING)
await hass.async_block_till_done()
assert hass.states.get("sensor.mock_name_port_1_poe_power")
await websocket_mock.reconnect()
assert (
hass.states.get("sensor.mock_name_port_1_poe_power").state != STATE_UNAVAILABLE
)
# Device gets disabled
device_1["disabled"] = True
@ -634,7 +636,10 @@ async def test_poe_port_switches(
async def test_wlan_client_sensors(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
websocket_mock,
) -> None:
"""Verify that WLAN client sensors are working as expected."""
wireless_client_1 = {
@ -720,13 +725,11 @@ async def test_wlan_client_sensors(
# Availability signalling
# Controller disconnects
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
await websocket_mock.disconnect()
assert hass.states.get("sensor.ssid_1").state == STATE_UNAVAILABLE
# Controller reconnects
mock_unifi_websocket(state=WebsocketState.RUNNING)
await hass.async_block_till_done()
await websocket_mock.reconnect()
assert hass.states.get("sensor.ssid_1").state == "0"
# WLAN gets disabled
@ -837,7 +840,6 @@ async def test_device_uptime(
now = datetime(2021, 1, 1, 1, 1, 0, tzinfo=dt_util.UTC)
with patch("homeassistant.util.dt.now", return_value=now):
await setup_unifi_integration(hass, aioclient_mock, devices_response=[device])
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 1
assert hass.states.get("sensor.device_uptime").state == "2021-01-01T01:00:00+00:00"
@ -854,7 +856,6 @@ async def test_device_uptime(
now = datetime(2021, 1, 1, 1, 1, 4, tzinfo=dt_util.UTC)
with patch("homeassistant.util.dt.now", return_value=now):
mock_unifi_websocket(message=MessageKey.DEVICE, data=device)
await hass.async_block_till_done()
assert hass.states.get("sensor.device_uptime").state == "2021-01-01T01:00:00+00:00"
@ -865,7 +866,6 @@ async def test_device_uptime(
now = datetime(2021, 2, 1, 1, 1, 0, tzinfo=dt_util.UTC)
with patch("homeassistant.util.dt.now", return_value=now):
mock_unifi_websocket(message=MessageKey.DEVICE, data=device)
await hass.async_block_till_done()
assert hass.states.get("sensor.device_uptime").state == "2021-02-01T01:00:00+00:00"
@ -908,5 +908,4 @@ async def test_device_temperature(
# Verify new event change temperature
device["general_temperature"] = 60
mock_unifi_websocket(message=MessageKey.DEVICE, data=device)
await hass.async_block_till_done()
assert hass.states.get("sensor.device_temperature").state == "60"

52
tests/components/unifi/test_switch.py
View File

@ -3,7 +3,6 @@ from copy import deepcopy
from datetime import timedelta
from aiounifi.models.message import MessageKey
from aiounifi.websocket import WebsocketState
import pytest
from homeassistant.components.switch import (
@ -1001,7 +1000,10 @@ async def test_block_switches(
async def test_dpi_switches(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
websocket_mock,
) -> None:
"""Test the update_items function with some clients."""
await setup_unifi_integration(
@ -1026,13 +1028,11 @@ async def test_dpi_switches(
# Availability signalling
# Controller disconnects
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
await websocket_mock.disconnect()
assert hass.states.get("switch.block_media_streaming").state == STATE_UNAVAILABLE
# Controller reconnects
mock_unifi_websocket(state=WebsocketState.RUNNING)
await hass.async_block_till_done()
await websocket_mock.reconnect()
assert hass.states.get("switch.block_media_streaming").state == STATE_OFF
# Remove app
@ -1128,6 +1128,7 @@ async def test_outlet_switches(
hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
websocket_mock,
entity_id: str,
test_data: any,
outlet_index: int,
@ -1192,13 +1193,11 @@ async def test_outlet_switches(
# Availability signalling
# Controller disconnects
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
await websocket_mock.disconnect()
assert hass.states.get(f"switch.{entity_id}").state == STATE_UNAVAILABLE
# Controller reconnects
mock_unifi_websocket(state=WebsocketState.RUNNING)
await hass.async_block_till_done()
await websocket_mock.reconnect()
assert hass.states.get(f"switch.{entity_id}").state == STATE_OFF
# Device gets disabled
@ -1320,7 +1319,10 @@ async def test_option_remove_switches(
async def test_poe_port_switches(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
websocket_mock,
) -> None:
"""Test the update_items function with some clients."""
config_entry = await setup_unifi_integration(
@ -1408,13 +1410,11 @@ async def test_poe_port_switches(
# Availability signalling
# Controller disconnects
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
await websocket_mock.disconnect()
assert hass.states.get("switch.mock_name_port_1_poe").state == STATE_UNAVAILABLE
# Controller reconnects
mock_unifi_websocket(state=WebsocketState.RUNNING)
await hass.async_block_till_done()
await websocket_mock.reconnect()
assert hass.states.get("switch.mock_name_port_1_poe").state == STATE_OFF
# Device gets disabled
@ -1431,7 +1431,10 @@ async def test_poe_port_switches(
async def test_wlan_switches(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
websocket_mock,
) -> None:
"""Test control of UniFi WLAN availability."""
config_entry = await setup_unifi_integration(
@ -1488,18 +1491,19 @@ async def test_wlan_switches(
# Availability signalling
# Controller disconnects
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
await websocket_mock.disconnect()
assert hass.states.get("switch.ssid_1").state == STATE_UNAVAILABLE
# Controller reconnects
mock_unifi_websocket(state=WebsocketState.RUNNING)
await hass.async_block_till_done()
await websocket_mock.reconnect()
assert hass.states.get("switch.ssid_1").state == STATE_OFF
async def test_port_forwarding_switches(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker,
mock_unifi_websocket,
websocket_mock,
) -> None:
"""Test control of UniFi port forwarding."""
_data = {
@ -1570,13 +1574,11 @@ async def test_port_forwarding_switches(
# Availability signalling
# Controller disconnects
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
await websocket_mock.disconnect()
assert hass.states.get("switch.unifi_network_plex").state == STATE_UNAVAILABLE
# Controller reconnects
mock_unifi_websocket(state=WebsocketState.RUNNING)
await hass.async_block_till_done()
await websocket_mock.reconnect()
assert hass.states.get("switch.unifi_network_plex").state == STATE_OFF
# Remove entity on deleted message

17
tests/components/unifi/test_update.py
View File

@ -2,7 +2,6 @@
from copy import deepcopy
from aiounifi.models.message import MessageKey
from aiounifi.websocket import WebsocketState
from yarl import URL
from homeassistant.components.unifi.const import CONF_SITE_ID
@ -185,26 +184,18 @@ async def test_install(
async def test_controller_state_change(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock
) -> None:
"""Verify entities state reflect on controller becoming unavailable."""
await setup_unifi_integration(
hass,
aioclient_mock,
devices_response=[DEVICE_1],
)
await setup_unifi_integration(hass, aioclient_mock, devices_response=[DEVICE_1])
assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 1
assert hass.states.get("update.device_1").state == STATE_ON
# Controller unavailable
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
await websocket_mock.disconnect()
assert hass.states.get("update.device_1").state == STATE_UNAVAILABLE
# Controller available
mock_unifi_websocket(state=WebsocketState.RUNNING)
await hass.async_block_till_done()
await websocket_mock.reconnect()
assert hass.states.get("update.device_1").state == STATE_ON