Simplify mocking UniFi websocket (#80652)

* Simplify mocking UniFi websocket

* Fix one more
pull/80697/head
Robert Svensson 2022-10-20 19:37:13 +02:00 committed by GitHub
parent 5cb8749ae3
commit 92eaa539b6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 91 additions and 300 deletions

View File

@ -3,7 +3,8 @@ from __future__ import annotations
from unittest.mock import patch
from aiounifi.websocket import WebsocketSignal
from aiounifi.models.message import MessageKey
from aiounifi.websocket import WebsocketSignal, WebsocketState
import pytest
from homeassistant.helpers import device_registry as dr
@ -16,11 +17,24 @@ 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 = ""):
def make_websocket_call(
*,
message: MessageKey | None = None,
data: list[dict] | dict | None = None,
state: WebsocketState | None = None,
):
"""Generate a websocket call."""
if data:
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)

View File

@ -8,6 +8,7 @@ from unittest.mock import Mock, patch
import aiounifi
from aiounifi.models.event import EventKey
from aiounifi.models.message import MessageKey
from aiounifi.websocket import WebsocketState
import pytest
@ -397,21 +398,14 @@ async def test_wireless_client_event_calls_update_wireless_devices(
"homeassistant.components.unifi.controller.UniFiController.update_wireless_clients",
return_value=None,
) as wireless_clients_mock:
mock_unifi_websocket(
data={
"meta": {"rc": "ok", "message": "events"},
"data": [
{
event = {
"datetime": "2020-01-20T19:37:04Z",
"user": "00:00:00:00:00:01",
"key": EventKey.WIRELESS_CLIENT_CONNECTED.value,
"msg": "User[11:22:33:44:55:66] has connected to WLAN",
"time": 1579549024893,
}
],
},
)
mock_unifi_websocket(message=MessageKey.EVENT, data=event)
assert wireless_clients_mock.assert_called_once

View File

@ -57,12 +57,7 @@ async def test_tracked_wireless_clients(
# Updated timestamp marks client as home
client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
await hass.async_block_till_done()
assert hass.states.get("device_tracker.client").state == STATE_HOME
@ -78,12 +73,7 @@ async def test_tracked_wireless_clients(
# Same timestamp doesn't explicitly mark client as away
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
await hass.async_block_till_done()
assert hass.states.get("device_tracker.client").state == STATE_HOME
@ -158,12 +148,7 @@ async def test_tracked_clients(
# State change signalling works
client_1["last_seen"] += 1
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client_1],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=client_1)
await hass.async_block_till_done()
assert hass.states.get("device_tracker.client_1").state == STATE_HOME
@ -208,14 +193,8 @@ async def test_tracked_wireless_clients_event_source(
"msg": f'User{[client["mac"]]} has connected to AP[{client["ap_mac"]}] with SSID "{client["essid"]}" on "channel 44(na)"',
"_id": "5ea331fa30c49e00f90ddc1a",
}
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.EVENT.value},
"data": [event],
}
)
mock_unifi_websocket(message=MessageKey.EVENT, data=event)
await hass.async_block_till_done()
assert hass.states.get("device_tracker.client").state == STATE_HOME
# Disconnected event
@ -235,12 +214,7 @@ async def test_tracked_wireless_clients_event_source(
"msg": f'User{[client["mac"]]} disconnected from "{client["essid"]}" (7m 47s connected, 448.28K bytes, last AP[{client["ap_mac"]}])',
"_id": "5ea32ff730c49e00f90dca1a",
}
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.EVENT.value},
"data": [event],
}
)
mock_unifi_websocket(message=MessageKey.EVENT, data=event)
await hass.async_block_till_done()
assert hass.states.get("device_tracker.client").state == STATE_HOME
@ -258,14 +232,8 @@ async def test_tracked_wireless_clients_event_source(
# New data
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
await hass.async_block_till_done()
assert hass.states.get("device_tracker.client").state == STATE_HOME
# Disconnection event will be ignored
@ -285,12 +253,7 @@ async def test_tracked_wireless_clients_event_source(
"msg": f'User{[client["mac"]]} disconnected from "{client["essid"]}" (7m 47s connected, 448.28K bytes, last AP[{client["ap_mac"]}])',
"_id": "5ea32ff730c49e00f90dca1a",
}
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.EVENT.value},
"data": [event],
}
)
mock_unifi_websocket(message=MessageKey.EVENT, data=event)
await hass.async_block_till_done()
assert hass.states.get("device_tracker.client").state == STATE_HOME
@ -350,19 +313,8 @@ async def test_tracked_devices(
# State change signalling work
device_1["next_interval"] = 20
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.DEVICE.value},
"data": [device_1],
}
)
device_2["next_interval"] = 50
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.DEVICE.value},
"data": [device_2],
}
)
mock_unifi_websocket(message=MessageKey.DEVICE, data=[device_1, device_2])
await hass.async_block_till_done()
assert hass.states.get("device_tracker.device_1").state == STATE_HOME
@ -381,12 +333,7 @@ async def test_tracked_devices(
# Disabled device is unavailable
device_1["disabled"] = True
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.DEVICE.value},
"data": [device_1],
}
)
mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
await hass.async_block_till_done()
assert hass.states.get("device_tracker.device_1").state == STATE_UNAVAILABLE
@ -420,12 +367,7 @@ async def test_remove_clients(
# Remove client
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT_REMOVED.value},
"data": [client_1],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT_REMOVED, data=client_1)
await hass.async_block_till_done()
await hass.async_block_till_done()
@ -723,20 +665,11 @@ async def test_option_ssid_filter(
# Roams to SSID outside of filter
client["essid"] = "other_ssid"
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
# Data update while SSID filter is in effect shouldn't create the client
client_on_ssid2["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client_on_ssid2],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=client_on_ssid2)
await hass.async_block_till_done()
# SSID filter marks client as away
@ -754,18 +687,7 @@ async def test_option_ssid_filter(
client["last_seen"] += 1
client_on_ssid2["last_seen"] += 1
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client],
}
)
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client_on_ssid2],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=[client, client_on_ssid2])
await hass.async_block_till_done()
assert hass.states.get("device_tracker.client").state == STATE_HOME
@ -781,12 +703,7 @@ async def test_option_ssid_filter(
assert hass.states.get("device_tracker.client").state == STATE_NOT_HOME
client_on_ssid2["last_seen"] += 1
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client_on_ssid2],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=client_on_ssid2)
await hass.async_block_till_done()
# Client won't go away until after next update
@ -794,12 +711,7 @@ async def test_option_ssid_filter(
# Trigger update to get client marked as away
client_on_ssid2["last_seen"] += 1
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client_on_ssid2],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=client_on_ssid2)
await hass.async_block_till_done()
new_time = (
@ -843,12 +755,7 @@ async def test_wireless_client_go_wired_issue(
# Trigger wired bug
client["last_seen"] += 1
client["is_wired"] = True
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
await hass.async_block_till_done()
# Wired bug fix keeps client marked as wireless
@ -869,12 +776,7 @@ async def test_wireless_client_go_wired_issue(
# Try to mark client as connected
client["last_seen"] += 1
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
await hass.async_block_till_done()
# Make sure it don't go online again until wired bug disappears
@ -885,12 +787,7 @@ async def test_wireless_client_go_wired_issue(
# Make client wireless
client["last_seen"] += 1
client["is_wired"] = False
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
await hass.async_block_till_done()
# Client is no longer affected by wired bug and can be marked online
@ -929,12 +826,7 @@ async def test_option_ignore_wired_bug(
# Trigger wired bug
client["is_wired"] = True
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
await hass.async_block_till_done()
# Wired bug in effect
@ -955,12 +847,7 @@ async def test_option_ignore_wired_bug(
# Mark client as connected again
client["last_seen"] += 1
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
await hass.async_block_till_done()
# Ignoring wired bug allows client to go home again even while affected
@ -971,12 +858,7 @@ async def test_option_ignore_wired_bug(
# Make client wireless
client["last_seen"] += 1
client["is_wired"] = False
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [client],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
await hass.async_block_till_done()
# Client is wireless and still connected

View File

@ -87,12 +87,7 @@ async def test_bandwidth_sensors(hass, aioclient_mock, mock_unifi_websocket):
wireless_client["rx_bytes-r"] = 3456000000
wireless_client["tx_bytes-r"] = 7891000000
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [wireless_client],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=wireless_client)
await hass.async_block_till_done()
assert hass.states.get("sensor.wireless_client_rx").state == "3456.0"
@ -199,12 +194,7 @@ async def test_uptime_sensors(
uptime_client["uptime"] = event_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(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [uptime_client],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=uptime_client)
await hass.async_block_till_done()
assert hass.states.get("sensor.client1_uptime").state == "2021-01-01T01:00:00+00:00"
@ -215,12 +205,7 @@ async def test_uptime_sensors(
uptime_client["uptime"] = new_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(
data={
"meta": {"message": MessageKey.CLIENT.value},
"data": [uptime_client],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=uptime_client)
await hass.async_block_till_done()
assert hass.states.get("sensor.client1_uptime").state == "2021-02-01T01:00:00+00:00"
@ -308,12 +293,7 @@ async def test_remove_sensors(hass, aioclient_mock, mock_unifi_websocket):
# Remove wired client
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT_REMOVED.value},
"data": [wired_client],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT_REMOVED, data=wired_client)
await hass.async_block_till_done()
assert len(hass.states.async_all()) == 5

View File

@ -742,12 +742,7 @@ async def test_remove_switches(hass, aioclient_mock, mock_unifi_websocket):
assert hass.states.get("switch.block_client_2") is not None
assert hass.states.get("switch.block_media_streaming") is not None
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.CLIENT_REMOVED.value},
"data": [CLIENT_1, UNBLOCKED],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT_REMOVED, data=[CLIENT_1, UNBLOCKED])
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
@ -789,12 +784,7 @@ async def test_block_switches(hass, aioclient_mock, mock_unifi_websocket):
assert unblocked is not None
assert unblocked.state == "on"
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.EVENT.value},
"data": [EVENT_BLOCKED_CLIENT_UNBLOCKED],
}
)
mock_unifi_websocket(message=MessageKey.EVENT, data=EVENT_BLOCKED_CLIENT_UNBLOCKED)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
@ -802,12 +792,7 @@ async def test_block_switches(hass, aioclient_mock, mock_unifi_websocket):
assert blocked is not None
assert blocked.state == "on"
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.EVENT.value},
"data": [EVENT_BLOCKED_CLIENT_BLOCKED],
}
)
mock_unifi_websocket(message=MessageKey.EVENT, data=EVENT_BLOCKED_CLIENT_BLOCKED)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
@ -881,9 +866,6 @@ async def test_dpi_switches_add_second_app(hass, aioclient_mock, mock_unifi_webs
assert hass.states.get("switch.block_media_streaming").state == STATE_ON
second_app_event = {
"meta": {"rc": "ok", "message": "dpiapp:add"},
"data": [
{
"apps": [524292],
"blocked": False,
"cats": [],
@ -892,34 +874,25 @@ async def test_dpi_switches_add_second_app(hass, aioclient_mock, mock_unifi_webs
"site_id": "name",
"_id": "61783e89c1773a18c0c61f00",
}
],
}
mock_unifi_websocket(data=second_app_event)
mock_unifi_websocket(message=MessageKey.DPI_APP_ADDED, data=second_app_event)
await hass.async_block_till_done()
assert hass.states.get("switch.block_media_streaming").state == STATE_ON
add_second_app_to_group = {
"meta": {"rc": "ok", "message": "dpigroup:sync"},
"data": [
{
"_id": "5f976f4ae3c58f018ec7dff6",
"name": "Block Media Streaming",
"site_id": "name",
"dpiapp_ids": ["5f976f62e3c58f018ec7e17d", "61783e89c1773a18c0c61f00"],
}
],
}
mock_unifi_websocket(data=add_second_app_to_group)
mock_unifi_websocket(
message=MessageKey.DPI_GROUP_UPDATED, data=add_second_app_to_group
)
await hass.async_block_till_done()
assert hass.states.get("switch.block_media_streaming").state == STATE_OFF
second_app_event_enabled = {
"meta": {"rc": "ok", "message": "dpiapp:sync"},
"data": [
{
"apps": [524292],
"blocked": False,
"cats": [],
@ -928,9 +901,9 @@ async def test_dpi_switches_add_second_app(hass, aioclient_mock, mock_unifi_webs
"site_id": "name",
"_id": "61783e89c1773a18c0c61f00",
}
],
}
mock_unifi_websocket(data=second_app_event_enabled)
mock_unifi_websocket(
message=MessageKey.DPI_APP_UPDATED, data=second_app_event_enabled
)
await hass.async_block_till_done()
assert hass.states.get("switch.block_media_streaming").state == STATE_ON
@ -957,12 +930,7 @@ async def test_outlet_switches(hass, aioclient_mock, mock_unifi_websocket):
outlet_up1 = deepcopy(OUTLET_UP1)
outlet_up1["outlet_table"][0]["relay_state"] = True
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.DEVICE.value},
"data": [outlet_up1],
}
)
mock_unifi_websocket(message=MessageKey.DEVICE, data=outlet_up1)
await hass.async_block_till_done()
outlet = hass.states.get("switch.plug_outlet_1")
@ -1035,22 +1003,12 @@ async def test_new_client_discovered_on_block_control(
blocked = hass.states.get("switch.block_client_1")
assert blocked is None
mock_unifi_websocket(
data={
"meta": {"message": "sta:sync"},
"data": [BLOCKED],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=BLOCKED)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.EVENT.value},
"data": [EVENT_BLOCKED_CLIENT_CONNECTED],
}
)
mock_unifi_websocket(message=MessageKey.EVENT, data=EVENT_BLOCKED_CLIENT_CONNECTED)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
@ -1141,22 +1099,12 @@ async def test_new_client_discovered_on_poe_control(
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
mock_unifi_websocket(
data={
"meta": {"message": "sta:sync"},
"data": [CLIENT_2],
}
)
mock_unifi_websocket(message=MessageKey.CLIENT, data=CLIENT_2)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.EVENT.value},
"data": [EVENT_CLIENT_2_CONNECTED],
}
)
mock_unifi_websocket(message=MessageKey.EVENT, data=EVENT_CLIENT_2_CONNECTED)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
@ -1423,12 +1371,7 @@ async def test_poe_port_switches(hass, aioclient_mock, mock_unifi_websocket):
# Update state object
device_1 = deepcopy(DEVICE_1)
device_1["port_table"][0]["poe_mode"] = "off"
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.DEVICE.value},
"data": [device_1],
}
)
mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
await hass.async_block_till_done()
assert hass.states.get("switch.mock_name_port_1_poe").state == STATE_OFF
@ -1475,22 +1418,12 @@ async def test_poe_port_switches(hass, aioclient_mock, mock_unifi_websocket):
# Device gets disabled
device_1["disabled"] = True
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.DEVICE.value},
"data": [device_1],
}
)
mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
await hass.async_block_till_done()
assert hass.states.get("switch.mock_name_port_1_poe").state == STATE_UNAVAILABLE
# Device gets re-enabled
device_1["disabled"] = False
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.DEVICE.value},
"data": [device_1],
}
)
mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
await hass.async_block_till_done()
assert hass.states.get("switch.mock_name_port_1_poe").state == STATE_OFF

View File

@ -64,9 +64,7 @@ async def test_no_entities(hass, aioclient_mock):
assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 0
async def test_device_updates(
hass, aioclient_mock, mock_unifi_websocket, mock_device_registry
):
async def test_device_updates(hass, aioclient_mock, mock_unifi_websocket):
"""Test the update_items function with some devices."""
device_1 = deepcopy(DEVICE_1)
await setup_unifi_integration(
@ -102,12 +100,7 @@ async def test_device_updates(
# Simulate start of update
device_1["state"] = 4
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.DEVICE.value},
"data": [device_1],
}
)
mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
await hass.async_block_till_done()
device_1_state = hass.states.get("update.device_1")
@ -122,12 +115,7 @@ async def test_device_updates(
device_1["version"] = "4.3.17.11279"
device_1["upgradable"] = False
del device_1["upgrade_to_firmware"]
mock_unifi_websocket(
data={
"meta": {"message": MessageKey.DEVICE.value},
"data": [device_1],
}
)
mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
await hass.async_block_till_done()
device_1_state = hass.states.get("update.device_1")