Use remove_device helper

home-assistant · Apr 28, 2024 · 3d4f3d1 · 3d4f3d1
1 parent b9948b6
commit 3d4f3d1
Showing 1 changed file with 3 additions and 32 deletions.
diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py
@@ -175,25 +175,14 @@ async def test_remove_config_entry_device(
     )
 
     assert await async_setup_component(hass, "config", {})
-    await hass.async_block_till_done()
-
     ws_client = await hass_ws_client(hass)
 
     # Try to remove an active client from UI: not allowed
     device_entry = device_registry.async_get_device(
         connections={(dr.CONNECTION_NETWORK_MAC, client_1["mac"])}
     )
-    await ws_client.send_json(
-        {
-            "id": 1,
-            "type": "config/device_registry/remove_config_entry",
-            "config_entry_id": config_entry.entry_id,
-            "device_id": device_entry.id,
-        }
-    )
-    response = await ws_client.receive_json()
+    response = await ws_client.remove_device(device_entry.id, config_entry.entry_id)
     assert not response["success"]
-    await hass.async_block_till_done()
     assert device_registry.async_get_device(
         connections={(dr.CONNECTION_NETWORK_MAC, client_1["mac"])}
     )
@@ -202,17 +191,8 @@ async def test_remove_config_entry_device(
     device_entry = device_registry.async_get_device(
         connections={(dr.CONNECTION_NETWORK_MAC, device_1["mac"])}
     )
-    await ws_client.send_json(
-        {
-            "id": 2,
-            "type": "config/device_registry/remove_config_entry",
-            "config_entry_id": config_entry.entry_id,
-            "device_id": device_entry.id,
-        }
-    )
-    response = await ws_client.receive_json()
+    response = await ws_client.remove_device(device_entry.id, config_entry.entry_id)
     assert not response["success"]
-    await hass.async_block_till_done()
     assert device_registry.async_get_device(
         connections={(dr.CONNECTION_NETWORK_MAC, device_1["mac"])}
     )
@@ -225,17 +205,8 @@ async def test_remove_config_entry_device(
     device_entry = device_registry.async_get_device(
         connections={(dr.CONNECTION_NETWORK_MAC, client_2["mac"])}
     )
-    await ws_client.send_json(
-        {
-            "id": 3,
-            "type": "config/device_registry/remove_config_entry",
-            "config_entry_id": config_entry.entry_id,
-            "device_id": device_entry.id,
-        }
-    )
-    response = await ws_client.receive_json()
+    response = await ws_client.remove_device(device_entry.id, config_entry.entry_id)
     assert response["success"]
-    await hass.async_block_till_done()
     assert not device_registry.async_get_device(
         connections={(dr.CONNECTION_NETWORK_MAC, client_2["mac"])}
     )