Simplify mocking UniFi websocket (#80652)

* Simplify mocking UniFi websocket

* Fix one more

tests/components/unifi/conftest.py

@@ -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)

tests/components/unifi/test_controller.py

@@ -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(
+        event = {
-            data={
-                "meta": {"rc": "ok", "message": "events"},
-                "data": [
-                    {
             "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
 
 

tests/components/unifi/test_device_tracker.py

@@ -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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=client_1)
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.EVENT, data=event)
-        data={
-            "meta": {"message": MessageKey.EVENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.EVENT, data=event)
-        data={
-            "meta": {"message": MessageKey.EVENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.EVENT, data=event)
-        data={
-            "meta": {"message": MessageKey.EVENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.DEVICE, data=[device_1, device_2])
-        data={
-            "meta": {"message": MessageKey.DEVICE.value},
-            "data": [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(
+    mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
-        data={
-            "meta": {"message": MessageKey.DEVICE.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT_REMOVED, data=client_1)
-        data={
-            "meta": {"message": MessageKey.CLIENT_REMOVED.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
-        data={
+
-            "meta": {"message": MessageKey.CLIENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=client_on_ssid2)
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=[client, client_on_ssid2])
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "data": [client],
-        }
-    )
-    mock_unifi_websocket(
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "data": [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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=client_on_ssid2)
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=client_on_ssid2)
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=client)
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "data": [client],
-        }
-    )
     await hass.async_block_till_done()
 
     # Client is wireless and still connected

tests/components/unifi/test_sensor.py

@@ -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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=wireless_client)
-        data={
-            "meta": {"message": MessageKey.CLIENT.value},
-            "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(
+        mock_unifi_websocket(message=MessageKey.CLIENT, data=uptime_client)
-            data={
-                "meta": {"message": MessageKey.CLIENT.value},
-                "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(
+        mock_unifi_websocket(message=MessageKey.CLIENT, data=uptime_client)
-            data={
-                "meta": {"message": MessageKey.CLIENT.value},
-                "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT_REMOVED, data=wired_client)
-        data={
-            "meta": {"message": MessageKey.CLIENT_REMOVED.value},
-            "data": [wired_client],
-        }
-    )
     await hass.async_block_till_done()
 
     assert len(hass.states.async_all()) == 5

tests/components/unifi/test_switch.py

@@ -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(
+    mock_unifi_websocket(message=MessageKey.CLIENT_REMOVED, data=[CLIENT_1, UNBLOCKED])
-        data={
-            "meta": {"message": MessageKey.CLIENT_REMOVED.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.EVENT, data=EVENT_BLOCKED_CLIENT_UNBLOCKED)
-        data={
-            "meta": {"message": MessageKey.EVENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.EVENT, data=EVENT_BLOCKED_CLIENT_BLOCKED)
-        data={
-            "meta": {"message": MessageKey.EVENT.value},
-            "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(message=MessageKey.DPI_APP_ADDED, data=second_app_event)
-    }
-    mock_unifi_websocket(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(
-    }
+        message=MessageKey.DPI_GROUP_UPDATED, data=add_second_app_to_group
-
+    )
-    mock_unifi_websocket(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(
-    }
+        message=MessageKey.DPI_APP_UPDATED, data=second_app_event_enabled
-    mock_unifi_websocket(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(
+    mock_unifi_websocket(message=MessageKey.DEVICE, data=outlet_up1)
-        data={
-            "meta": {"message": MessageKey.DEVICE.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=BLOCKED)
-        data={
-            "meta": {"message": "sta:sync"},
-            "data": [BLOCKED],
-        }
-    )
     await hass.async_block_till_done()
 
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
 
-    mock_unifi_websocket(
+    mock_unifi_websocket(message=MessageKey.EVENT, data=EVENT_BLOCKED_CLIENT_CONNECTED)
-        data={
-            "meta": {"message": MessageKey.EVENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.CLIENT, data=CLIENT_2)
-        data={
-            "meta": {"message": "sta:sync"},
-            "data": [CLIENT_2],
-        }
-    )
     await hass.async_block_till_done()
 
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
 
-    mock_unifi_websocket(
+    mock_unifi_websocket(message=MessageKey.EVENT, data=EVENT_CLIENT_2_CONNECTED)
-        data={
-            "meta": {"message": MessageKey.EVENT.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
-        data={
-            "meta": {"message": MessageKey.DEVICE.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
-        data={
-            "meta": {"message": MessageKey.DEVICE.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
-        data={
-            "meta": {"message": MessageKey.DEVICE.value},
-            "data": [device_1],
-        }
-    )
     await hass.async_block_till_done()
     assert hass.states.get("switch.mock_name_port_1_poe").state == STATE_OFF

tests/components/unifi/test_update.py

@@ -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(
+async def test_device_updates(hass, aioclient_mock, mock_unifi_websocket):
-    hass, aioclient_mock, mock_unifi_websocket, mock_device_registry
-):
     """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(
+    mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
-        data={
-            "meta": {"message": MessageKey.DEVICE.value},
-            "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(
+    mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
-        data={
-            "meta": {"message": MessageKey.DEVICE.value},
-            "data": [device_1],
-        }
-    )
     await hass.async_block_till_done()
 
     device_1_state = hass.states.get("update.device_1")