Handle remove packets for UniFi Protect (#77337)

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/unifiprotect/button.py

@@ -129,7 +129,8 @@ async def async_setup_entry(
         async_add_entities(entities)
         _async_remove_adopt_button(hass, device)
 
-    async def _add_unadopted_device(device: ProtectAdoptableDeviceModel) -> None:
+    @callback
+    def _async_add_unadopted_device(device: ProtectAdoptableDeviceModel) -> None:
         if not device.can_adopt or not device.can_create(data.api.bootstrap.auth_user):
             _LOGGER.debug("Device is not adoptable: %s", device.id)
             return
@@ -147,7 +148,7 @@ async def async_setup_entry(
     )
     entry.async_on_unload(
         async_dispatcher_connect(
-            hass, _ufpd(entry, DISPATCH_ADD), _add_unadopted_device
+            hass, _ufpd(entry, DISPATCH_ADD), _async_add_unadopted_device
         )
     )
 

homeassistant/components/unifiprotect/data.py

@@ -21,6 +21,7 @@ from pyunifiprotect.exceptions import ClientError, NotAuthorized
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.event import async_track_time_interval
 
@@ -168,6 +169,18 @@ class ProtectData:
             _LOGGER.debug("New device detected: %s", device.id)
             async_dispatcher_send(self._hass, _ufpd(self._entry, DISPATCH_ADD), device)
 
+    @callback
+    def _async_remove_device(self, device: ProtectAdoptableDeviceModel) -> None:
+        registry = dr.async_get(self._hass)
+        device_entry = registry.async_get_device(
+            identifiers=set(), connections={(dr.CONNECTION_NETWORK_MAC, device.mac)}
+        )
+        if device_entry:
+            _LOGGER.debug("Device removed: %s", device.id)
+            registry.async_update_device(
+                device_entry.id, remove_config_entry_id=self._entry.entry_id
+            )
+
     @callback
     def _async_update_device(
         self, device: ProtectAdoptableDeviceModel | NVR, changed_data: dict[str, Any]
@@ -195,8 +208,9 @@ class ProtectData:
 
     @callback
     def _async_process_ws_message(self, message: WSSubscriptionMessage) -> None:
-        # removed packets are not processed yet
         if message.new_obj is None:
+            if isinstance(message.old_obj, ProtectAdoptableDeviceModel):
+                self._async_remove_device(message.old_obj)
             return
 
         obj = message.new_obj

homeassistant/components/unifiprotect/entity.py

@@ -41,7 +41,7 @@ def _async_device_entities(
     unadopted_descs: Sequence[ProtectRequiredKeysMixin],
     ufp_device: ProtectAdoptableDeviceModel | None = None,
 ) -> list[ProtectDeviceEntity]:
-    if len(descs) + len(unadopted_descs) == 0:
+    if not descs and not unadopted_descs:
         return []
 
     entities: list[ProtectDeviceEntity] = []
@@ -162,7 +162,7 @@ def async_all_device_entities(
     elif ufp_device.model == ModelType.CHIME:
         descs = chime_descs
 
-    if len(descs) + len(unadopted_descs) == 0 or ufp_device.model is None:
+    if not descs and not unadopted_descs or ufp_device.model is None:
         return []
     return _async_device_entities(
         data, klass, ufp_device.model, descs, unadopted_descs, ufp_device

tests/components/unifiprotect/test_binary_sensor.py

@@ -51,7 +51,7 @@ async def test_binary_sensor_camera_remove(
     ufp.api.bootstrap.nvr.system_info.ustorage = None
     await init_entry(hass, ufp, [doorbell, unadopted_camera])
     assert_entity_counts(hass, Platform.BINARY_SENSOR, 3, 3)
-    await remove_entities(hass, [doorbell, unadopted_camera])
+    await remove_entities(hass, ufp, [doorbell, unadopted_camera])
     assert_entity_counts(hass, Platform.BINARY_SENSOR, 0, 0)
     await adopt_devices(hass, ufp, [doorbell, unadopted_camera])
     assert_entity_counts(hass, Platform.BINARY_SENSOR, 3, 3)
@@ -65,7 +65,7 @@ async def test_binary_sensor_light_remove(
     ufp.api.bootstrap.nvr.system_info.ustorage = None
     await init_entry(hass, ufp, [light])
     assert_entity_counts(hass, Platform.BINARY_SENSOR, 2, 2)
-    await remove_entities(hass, [light])
+    await remove_entities(hass, ufp, [light])
     assert_entity_counts(hass, Platform.BINARY_SENSOR, 0, 0)
     await adopt_devices(hass, ufp, [light])
     assert_entity_counts(hass, Platform.BINARY_SENSOR, 2, 2)
@@ -79,7 +79,7 @@ async def test_binary_sensor_sensor_remove(
     ufp.api.bootstrap.nvr.system_info.ustorage = None
     await init_entry(hass, ufp, [sensor_all])
     assert_entity_counts(hass, Platform.BINARY_SENSOR, 4, 4)
-    await remove_entities(hass, [sensor_all])
+    await remove_entities(hass, ufp, [sensor_all])
     assert_entity_counts(hass, Platform.BINARY_SENSOR, 0, 0)
     await adopt_devices(hass, ufp, [sensor_all])
     assert_entity_counts(hass, Platform.BINARY_SENSOR, 4, 4)

tests/components/unifiprotect/test_button.py

@@ -28,7 +28,7 @@ async def test_button_chime_remove(
 
     await init_entry(hass, ufp, [chime])
     assert_entity_counts(hass, Platform.BUTTON, 4, 2)
-    await remove_entities(hass, [chime])
+    await remove_entities(hass, ufp, [chime])
     assert_entity_counts(hass, Platform.BUTTON, 0, 0)
     await adopt_devices(hass, ufp, [chime])
     assert_entity_counts(hass, Platform.BUTTON, 4, 2)

tests/components/unifiprotect/test_camera.py

@@ -279,7 +279,7 @@ async def test_adopt(hass: HomeAssistant, ufp: MockUFPFixture, camera: ProtectCa
     await init_entry(hass, ufp, [camera1])
     assert_entity_counts(hass, Platform.CAMERA, 0, 0)
 
-    await remove_entities(hass, [camera1])
+    await remove_entities(hass, ufp, [camera1])
     assert_entity_counts(hass, Platform.CAMERA, 0, 0)
     camera1.channels = []
     await adopt_devices(hass, ufp, [camera1])
@@ -296,7 +296,7 @@ async def test_adopt(hass: HomeAssistant, ufp: MockUFPFixture, camera: ProtectCa
     await hass.async_block_till_done()
     assert_entity_counts(hass, Platform.CAMERA, 2, 1)
 
-    await remove_entities(hass, [camera1])
+    await remove_entities(hass, ufp, [camera1])
     assert_entity_counts(hass, Platform.CAMERA, 0, 0)
     await adopt_devices(hass, ufp, [camera1])
     assert_entity_counts(hass, Platform.CAMERA, 2, 1)

tests/components/unifiprotect/test_light.py

@@ -33,7 +33,7 @@ async def test_light_remove(hass: HomeAssistant, ufp: MockUFPFixture, light: Lig
 
     await init_entry(hass, ufp, [light])
     assert_entity_counts(hass, Platform.LIGHT, 1, 1)
-    await remove_entities(hass, [light])
+    await remove_entities(hass, ufp, [light])
     assert_entity_counts(hass, Platform.LIGHT, 0, 0)
     await adopt_devices(hass, ufp, [light])
     assert_entity_counts(hass, Platform.LIGHT, 1, 1)

tests/components/unifiprotect/test_lock.py

@@ -37,7 +37,7 @@ async def test_lock_remove(
 
     await init_entry(hass, ufp, [doorlock])
     assert_entity_counts(hass, Platform.LOCK, 1, 1)
-    await remove_entities(hass, [doorlock])
+    await remove_entities(hass, ufp, [doorlock])
     assert_entity_counts(hass, Platform.LOCK, 0, 0)
     await adopt_devices(hass, ufp, [doorlock])
     assert_entity_counts(hass, Platform.LOCK, 1, 1)

tests/components/unifiprotect/test_media_player.py

@@ -41,7 +41,7 @@ async def test_media_player_camera_remove(
 
     await init_entry(hass, ufp, [doorbell])
     assert_entity_counts(hass, Platform.MEDIA_PLAYER, 1, 1)
-    await remove_entities(hass, [doorbell])
+    await remove_entities(hass, ufp, [doorbell])
     assert_entity_counts(hass, Platform.MEDIA_PLAYER, 0, 0)
     await adopt_devices(hass, ufp, [doorbell])
     assert_entity_counts(hass, Platform.MEDIA_PLAYER, 1, 1)

tests/components/unifiprotect/test_number.py

@@ -36,7 +36,7 @@ async def test_number_sensor_camera_remove(
 
     await init_entry(hass, ufp, [camera, unadopted_camera])
     assert_entity_counts(hass, Platform.NUMBER, 3, 3)
-    await remove_entities(hass, [camera, unadopted_camera])
+    await remove_entities(hass, ufp, [camera, unadopted_camera])
     assert_entity_counts(hass, Platform.NUMBER, 0, 0)
     await adopt_devices(hass, ufp, [camera, unadopted_camera])
     assert_entity_counts(hass, Platform.NUMBER, 3, 3)
@@ -49,7 +49,7 @@ async def test_number_sensor_light_remove(
 
     await init_entry(hass, ufp, [light])
     assert_entity_counts(hass, Platform.NUMBER, 2, 2)
-    await remove_entities(hass, [light])
+    await remove_entities(hass, ufp, [light])
     assert_entity_counts(hass, Platform.NUMBER, 0, 0)
     await adopt_devices(hass, ufp, [light])
     assert_entity_counts(hass, Platform.NUMBER, 2, 2)
@@ -62,7 +62,7 @@ async def test_number_lock_remove(
 
     await init_entry(hass, ufp, [doorlock])
     assert_entity_counts(hass, Platform.NUMBER, 1, 1)
-    await remove_entities(hass, [doorlock])
+    await remove_entities(hass, ufp, [doorlock])
     assert_entity_counts(hass, Platform.NUMBER, 0, 0)
     await adopt_devices(hass, ufp, [doorlock])
     assert_entity_counts(hass, Platform.NUMBER, 1, 1)

tests/components/unifiprotect/test_select.py

@@ -57,7 +57,7 @@ async def test_select_camera_remove(
     ufp.api.bootstrap.nvr.system_info.ustorage = None
     await init_entry(hass, ufp, [doorbell, unadopted_camera])
     assert_entity_counts(hass, Platform.SELECT, 4, 4)
-    await remove_entities(hass, [doorbell, unadopted_camera])
+    await remove_entities(hass, ufp, [doorbell, unadopted_camera])
     assert_entity_counts(hass, Platform.SELECT, 0, 0)
     await adopt_devices(hass, ufp, [doorbell, unadopted_camera])
     assert_entity_counts(hass, Platform.SELECT, 4, 4)
@@ -71,7 +71,7 @@ async def test_select_light_remove(
     ufp.api.bootstrap.nvr.system_info.ustorage = None
     await init_entry(hass, ufp, [light])
     assert_entity_counts(hass, Platform.SELECT, 2, 2)
-    await remove_entities(hass, [light])
+    await remove_entities(hass, ufp, [light])
     assert_entity_counts(hass, Platform.SELECT, 0, 0)
     await adopt_devices(hass, ufp, [light])
     assert_entity_counts(hass, Platform.SELECT, 2, 2)
@@ -85,7 +85,7 @@ async def test_select_viewer_remove(
     ufp.api.bootstrap.nvr.system_info.ustorage = None
     await init_entry(hass, ufp, [viewer])
     assert_entity_counts(hass, Platform.SELECT, 1, 1)
-    await remove_entities(hass, [viewer])
+    await remove_entities(hass, ufp, [viewer])
     assert_entity_counts(hass, Platform.SELECT, 0, 0)
     await adopt_devices(hass, ufp, [viewer])
     assert_entity_counts(hass, Platform.SELECT, 1, 1)

tests/components/unifiprotect/test_sensor.py

@@ -63,7 +63,7 @@ async def test_sensor_camera_remove(
     ufp.api.bootstrap.nvr.system_info.ustorage = None
     await init_entry(hass, ufp, [doorbell, unadopted_camera])
     assert_entity_counts(hass, Platform.SENSOR, 25, 13)
-    await remove_entities(hass, [doorbell, unadopted_camera])
+    await remove_entities(hass, ufp, [doorbell, unadopted_camera])
     assert_entity_counts(hass, Platform.SENSOR, 12, 9)
     await adopt_devices(hass, ufp, [doorbell, unadopted_camera])
     assert_entity_counts(hass, Platform.SENSOR, 25, 13)
@@ -77,7 +77,7 @@ async def test_sensor_sensor_remove(
     ufp.api.bootstrap.nvr.system_info.ustorage = None
     await init_entry(hass, ufp, [sensor_all])
     assert_entity_counts(hass, Platform.SENSOR, 22, 14)
-    await remove_entities(hass, [sensor_all])
+    await remove_entities(hass, ufp, [sensor_all])
     assert_entity_counts(hass, Platform.SENSOR, 12, 9)
     await adopt_devices(hass, ufp, [sensor_all])
     assert_entity_counts(hass, Platform.SENSOR, 22, 14)

tests/components/unifiprotect/test_switch.py

@@ -50,7 +50,7 @@ async def test_switch_camera_remove(
     ufp.api.bootstrap.nvr.system_info.ustorage = None
     await init_entry(hass, ufp, [doorbell, unadopted_camera])
     assert_entity_counts(hass, Platform.SWITCH, 13, 12)
-    await remove_entities(hass, [doorbell, unadopted_camera])
+    await remove_entities(hass, ufp, [doorbell, unadopted_camera])
     assert_entity_counts(hass, Platform.SWITCH, 0, 0)
     await adopt_devices(hass, ufp, [doorbell, unadopted_camera])
     assert_entity_counts(hass, Platform.SWITCH, 13, 12)
@@ -64,7 +64,7 @@ async def test_switch_light_remove(
     ufp.api.bootstrap.nvr.system_info.ustorage = None
     await init_entry(hass, ufp, [light])
     assert_entity_counts(hass, Platform.SWITCH, 2, 1)
-    await remove_entities(hass, [light])
+    await remove_entities(hass, ufp, [light])
     assert_entity_counts(hass, Platform.SWITCH, 0, 0)
     await adopt_devices(hass, ufp, [light])
     assert_entity_counts(hass, Platform.SWITCH, 2, 1)

tests/components/unifiprotect/utils.py

@@ -23,7 +23,7 @@ from pyunifiprotect.test_util.anonymize import random_hex
 
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant, split_entity_id
-from homeassistant.helpers import device_registry as dr, entity_registry as er
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity import EntityDescription
 import homeassistant.util.dt as dt_util
 
@@ -89,7 +89,6 @@ def assert_entity_counts(
         e for e in entity_registry.entities if split_entity_id(e)[0] == platform.value
     ]
 
-    print(len(entities), total)
     assert len(entities) == total
     assert len(hass.states.async_all(platform.value)) == enabled
 
@@ -176,28 +175,25 @@ async def init_entry(
 
 async def remove_entities(
     hass: HomeAssistant,
+    ufp: MockUFPFixture,
     ufp_devices: list[ProtectAdoptableDeviceModel],
 ) -> None:
     """Remove all entities for given Protect devices."""
 
-    entity_registry = er.async_get(hass)
-    device_registry = dr.async_get(hass)
-
     for ufp_device in ufp_devices:
         if not ufp_device.is_adopted_by_us:
             continue
 
-        name = ufp_device.display_name.replace(" ", "_").lower()
+        devices = getattr(ufp.api.bootstrap, f"{ufp_device.model.value}s")
-        entity = entity_registry.async_get(f"{Platform.SENSOR}.{name}_uptime")
+        del devices[ufp_device.id]
-        assert entity is not None
 
-        device_id = entity.device_id
+        mock_msg = Mock()
-        for reg in list(entity_registry.entities.values()):
+        mock_msg.changed_data = {}
-            if reg.device_id == device_id:
+        mock_msg.old_obj = ufp_device
-                entity_registry.async_remove(reg.entity_id)
+        mock_msg.new_obj = None
-        device_registry.async_remove_device(device_id)
+        ufp.ws_msg(mock_msg)
 
-    await hass.async_block_till_done()
+    await time_changed(hass, 30)
 
 
 async def adopt_devices(
@@ -214,6 +210,9 @@ async def adopt_devices(
             ufp_device.is_adopted_by_other = False
             ufp_device.can_adopt = False
 
+        devices = getattr(ufp.api.bootstrap, f"{ufp_device.model.value}s")
+        devices[ufp_device.id] = ufp_device
+
         mock_msg = Mock()
         mock_msg.changed_data = {}
         mock_msg.new_obj = Event(