Add adopt/unadopt flows for UniFi Protect devices (#76524)

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

homeassistant/components/unifiprotect/button.py

@@ -2,9 +2,10 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+import logging
 from typing import Final
 
-from pyunifiprotect.data import ProtectAdoptableDeviceModel
+from pyunifiprotect.data import ProtectAdoptableDeviceModel, ProtectModelWithId
 
 from homeassistant.components.button import (
     ButtonDeviceClass,
@@ -12,16 +13,20 @@ from homeassistant.components.button import (
     ButtonEntityDescription,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.core import HomeAssistant
+from homeassistant.const import Platform
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DISPATCH_ADOPT, DOMAIN
+from .const import DEVICES_THAT_ADOPT, DISPATCH_ADD, DISPATCH_ADOPT, DOMAIN
 from .data import ProtectData
 from .entity import ProtectDeviceEntity, async_all_device_entities
 from .models import PermRequired, ProtectSetableKeysMixin, T
 from .utils import async_dispatch_id as _ufpd
 
+_LOGGER = logging.getLogger(__name__)
+
 
 @dataclass
 class ProtectButtonEntityDescription(
@@ -33,6 +38,7 @@ class ProtectButtonEntityDescription(
 
 
 DEVICE_CLASS_CHIME_BUTTON: Final = "unifiprotect__chime_button"
+KEY_ADOPT = "adopt"
 
 
 ALL_DEVICE_BUTTONS: tuple[ProtectButtonEntityDescription, ...] = (
@@ -44,6 +50,21 @@ ALL_DEVICE_BUTTONS: tuple[ProtectButtonEntityDescription, ...] = (
         ufp_press="reboot",
         ufp_perm=PermRequired.WRITE,
     ),
+    ProtectButtonEntityDescription(
+        key="unadopt",
+        entity_registry_enabled_default=False,
+        name="Unadopt Device",
+        icon="mdi:delete",
+        ufp_press="unadopt",
+        ufp_perm=PermRequired.DELETE,
+    ),
+)
+
+ADOPT_BUTTON = ProtectButtonEntityDescription[ProtectAdoptableDeviceModel](
+    key=KEY_ADOPT,
+    name="Adopt Device",
+    icon="mdi:plus-circle",
+    ufp_press="adopt",
 )
 
 SENSOR_BUTTONS: tuple[ProtectButtonEntityDescription, ...] = (
@@ -73,6 +94,20 @@ CHIME_BUTTONS: tuple[ProtectButtonEntityDescription, ...] = (
 )
 
 
+@callback
+def _async_remove_adopt_button(
+    hass: HomeAssistant, device: ProtectAdoptableDeviceModel
+) -> None:
+
+    entity_registry = er.async_get(hass)
+    if device.is_adopted_by_us and (
+        entity_id := entity_registry.async_get_entity_id(
+            Platform.BUTTON, DOMAIN, f"{device.mac}_adopt"
+        )
+    ):
+        entity_registry.async_remove(entity_id)
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     entry: ConfigEntry,
@@ -86,25 +121,49 @@ async def async_setup_entry(
             data,
             ProtectButton,
             all_descs=ALL_DEVICE_BUTTONS,
+            unadopted_descs=[ADOPT_BUTTON],
             chime_descs=CHIME_BUTTONS,
             sense_descs=SENSOR_BUTTONS,
             ufp_device=device,
         )
         async_add_entities(entities)
+        _async_remove_adopt_button(hass, device)
+
+    async def _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
+
+        entities = async_all_device_entities(
+            data,
+            ProtectButton,
+            unadopted_descs=[ADOPT_BUTTON],
+            ufp_device=device,
+        )
+        async_add_entities(entities)
 
     entry.async_on_unload(
         async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device)
     )
+    entry.async_on_unload(
+        async_dispatcher_connect(
+            hass, _ufpd(entry, DISPATCH_ADD), _add_unadopted_device
+        )
+    )
 
     entities: list[ProtectDeviceEntity] = async_all_device_entities(
         data,
         ProtectButton,
         all_descs=ALL_DEVICE_BUTTONS,
+        unadopted_descs=[ADOPT_BUTTON],
         chime_descs=CHIME_BUTTONS,
         sense_descs=SENSOR_BUTTONS,
     )
     async_add_entities(entities)
 
+    for device in data.get_by_types(DEVICES_THAT_ADOPT):
+        _async_remove_adopt_button(hass, device)
+
 
 class ProtectButton(ProtectDeviceEntity, ButtonEntity):
     """A Ubiquiti UniFi Protect Reboot button."""
@@ -121,6 +180,15 @@ class ProtectButton(ProtectDeviceEntity, ButtonEntity):
         super().__init__(data, device, description)
         self._attr_name = f"{self.device.display_name} {self.entity_description.name}"
 
+    @callback
+    def _async_update_device_from_protect(self, device: ProtectModelWithId) -> None:
+        super()._async_update_device_from_protect(device)
+
+        if self.entity_description.key == KEY_ADOPT:
+            self._attr_available = self.device.can_adopt and self.device.can_create(
+                self.data.api.bootstrap.auth_user
+            )
+
     async def async_press(self) -> None:
         """Press the button."""
 

homeassistant/components/unifiprotect/const.py

@@ -63,5 +63,6 @@ PLATFORMS = [
     Platform.SWITCH,
 ]
 
+DISPATCH_ADD = "add_device"
 DISPATCH_ADOPT = "adopt_device"
 DISPATCH_CHANNELS = "new_camera_channels"

homeassistant/components/unifiprotect/data.py

@@ -29,6 +29,7 @@ from .const import (
     CONF_MAX_MEDIA,
     DEFAULT_MAX_MEDIA,
     DEVICES_THAT_ADOPT,
+    DISPATCH_ADD,
     DISPATCH_ADOPT,
     DISPATCH_CHANNELS,
     DOMAIN,
@@ -157,28 +158,33 @@ class ProtectData:
         self._pending_camera_ids.add(camera_id)
 
     @callback
-    def _async_process_ws_message(self, message: WSSubscriptionMessage) -> None:
+    def _async_add_device(self, device: ProtectAdoptableDeviceModel) -> None:
-        # removed packets are not processed yet
+        if device.is_adopted_by_us:
-        if message.new_obj is None or not getattr(
+            _LOGGER.debug("Device adopted: %s", device.id)
-            message.new_obj, "is_adopted_by_us", True
-        ):
-            return
-
-        obj = message.new_obj
-        if isinstance(obj, (ProtectAdoptableDeviceModel, NVR)):
-            self._async_signal_device_update(obj)
-            if (
-                obj.model == ModelType.CAMERA
-                and obj.id in self._pending_camera_ids
-                and "channels" in message.changed_data
-            ):
-                self._pending_camera_ids.remove(obj.id)
             async_dispatcher_send(
-                    self._hass, _ufpd(self._entry, DISPATCH_CHANNELS), obj
+                self._hass, _ufpd(self._entry, DISPATCH_ADOPT), device
+            )
+        else:
+            _LOGGER.debug("New device detected: %s", device.id)
+            async_dispatcher_send(self._hass, _ufpd(self._entry, DISPATCH_ADD), device)
+
+    @callback
+    def _async_update_device(
+        self, device: ProtectAdoptableDeviceModel | NVR, changed_data: dict[str, Any]
+    ) -> None:
+        self._async_signal_device_update(device)
+        if (
+            device.model == ModelType.CAMERA
+            and device.id in self._pending_camera_ids
+            and "channels" in changed_data
+        ):
+            self._pending_camera_ids.remove(device.id)
+            async_dispatcher_send(
+                self._hass, _ufpd(self._entry, DISPATCH_CHANNELS), device
             )
 
         # trigger update for all Cameras with LCD screens when NVR Doorbell settings updates
-            if "doorbell_settings" in message.changed_data:
+        if "doorbell_settings" in changed_data:
             _LOGGER.debug(
                 "Doorbell messages updated. Updating devices with LCD screens"
             )
@@ -186,6 +192,20 @@ class ProtectData:
             for camera in self.api.bootstrap.cameras.values():
                 if camera.feature_flags.has_lcd_screen:
                     self._async_signal_device_update(camera)
+
+    @callback
+    def _async_process_ws_message(self, message: WSSubscriptionMessage) -> None:
+        # removed packets are not processed yet
+        if message.new_obj is None:
+            return
+
+        obj = message.new_obj
+        if isinstance(obj, (ProtectAdoptableDeviceModel, NVR)):
+            if message.old_obj is None and isinstance(obj, ProtectAdoptableDeviceModel):
+                self._async_add_device(obj)
+            elif getattr(obj, "is_adopted_by_us", True):
+                self._async_update_device(obj, message.changed_data)
+
         # trigger updates for camera that the event references
         elif isinstance(obj, Event):
             if obj.type == EventType.DEVICE_ADOPTED:
@@ -194,10 +214,7 @@ class ProtectData:
                         obj.metadata.device_id
                     )
                     if device is not None:
-                        _LOGGER.debug("New device detected: %s", device.id)
+                        self._async_add_device(device)
-                        async_dispatcher_send(
-                            self._hass, _ufpd(self._entry, DISPATCH_ADOPT), device
-                        )
             elif obj.camera is not None:
                 self._async_signal_device_update(obj.camera)
             elif obj.light is not None:

homeassistant/components/unifiprotect/entity.py

@@ -38,9 +38,10 @@ def _async_device_entities(
     klass: type[ProtectDeviceEntity],
     model_type: ModelType,
     descs: Sequence[ProtectRequiredKeysMixin],
+    unadopted_descs: Sequence[ProtectRequiredKeysMixin],
     ufp_device: ProtectAdoptableDeviceModel | None = None,
 ) -> list[ProtectDeviceEntity]:
-    if len(descs) == 0:
+    if len(descs) + len(unadopted_descs) == 0:
         return []
 
     entities: list[ProtectDeviceEntity] = []
@@ -48,17 +49,36 @@ def _async_device_entities(
         [ufp_device] if ufp_device is not None else data.get_by_types({model_type})
     )
     for device in devices:
+        assert isinstance(device, (Camera, Light, Sensor, Viewer, Doorlock, Chime))
         if not device.is_adopted_by_us:
+            for description in unadopted_descs:
+                entities.append(
+                    klass(
+                        data,
+                        device=device,
+                        description=description,
+                    )
+                )
+                _LOGGER.debug(
+                    "Adding %s entity %s for %s",
+                    klass.__name__,
+                    description.name,
+                    device.display_name,
+                )
             continue
 
-        assert isinstance(device, (Camera, Light, Sensor, Viewer, Doorlock, Chime))
+        can_write = device.can_write(data.api.bootstrap.auth_user)
         for description in descs:
             if description.ufp_perm is not None:
-                can_write = device.can_write(data.api.bootstrap.auth_user)
                 if description.ufp_perm == PermRequired.WRITE and not can_write:
                     continue
                 if description.ufp_perm == PermRequired.NO_WRITE and can_write:
                     continue
+                if (
+                    description.ufp_perm == PermRequired.DELETE
+                    and not device.can_delete(data.api.bootstrap.auth_user)
+                ):
+                    continue
 
             if description.ufp_required_field:
                 required_field = get_nested_attr(device, description.ufp_required_field)
@@ -93,10 +113,12 @@ def async_all_device_entities(
     lock_descs: Sequence[ProtectRequiredKeysMixin] | None = None,
     chime_descs: Sequence[ProtectRequiredKeysMixin] | None = None,
     all_descs: Sequence[ProtectRequiredKeysMixin] | None = None,
+    unadopted_descs: Sequence[ProtectRequiredKeysMixin] | None = None,
     ufp_device: ProtectAdoptableDeviceModel | None = None,
 ) -> list[ProtectDeviceEntity]:
     """Generate a list of all the device entities."""
     all_descs = list(all_descs or [])
+    unadopted_descs = list(unadopted_descs or [])
     camera_descs = list(camera_descs or []) + all_descs
     light_descs = list(light_descs or []) + all_descs
     sense_descs = list(sense_descs or []) + all_descs
@@ -106,12 +128,24 @@ def async_all_device_entities(
 
     if ufp_device is None:
         return (
-            _async_device_entities(data, klass, ModelType.CAMERA, camera_descs)
+            _async_device_entities(
-            + _async_device_entities(data, klass, ModelType.LIGHT, light_descs)
+                data, klass, ModelType.CAMERA, camera_descs, unadopted_descs
-            + _async_device_entities(data, klass, ModelType.SENSOR, sense_descs)
+            )
-            + _async_device_entities(data, klass, ModelType.VIEWPORT, viewer_descs)
+            + _async_device_entities(
-            + _async_device_entities(data, klass, ModelType.DOORLOCK, lock_descs)
+                data, klass, ModelType.LIGHT, light_descs, unadopted_descs
-            + _async_device_entities(data, klass, ModelType.CHIME, chime_descs)
+            )
+            + _async_device_entities(
+                data, klass, ModelType.SENSOR, sense_descs, unadopted_descs
+            )
+            + _async_device_entities(
+                data, klass, ModelType.VIEWPORT, viewer_descs, unadopted_descs
+            )
+            + _async_device_entities(
+                data, klass, ModelType.DOORLOCK, lock_descs, unadopted_descs
+            )
+            + _async_device_entities(
+                data, klass, ModelType.CHIME, chime_descs, unadopted_descs
+            )
         )
 
     descs = []
@@ -128,9 +162,11 @@ def async_all_device_entities(
     elif ufp_device.model == ModelType.CHIME:
         descs = chime_descs
 
-    if len(descs) == 0 or ufp_device.model is None:
+    if len(descs) + len(unadopted_descs) == 0 or ufp_device.model is None:
         return []
-    return _async_device_entities(data, klass, ufp_device.model, descs, ufp_device)
+    return _async_device_entities(
+        data, klass, ufp_device.model, descs, unadopted_descs, ufp_device
+    )
 
 
 class ProtectDeviceEntity(Entity):
@@ -190,8 +226,9 @@ class ProtectDeviceEntity(Entity):
             assert isinstance(device, ProtectAdoptableDeviceModel)
             self.device = device
 
-        is_connected = (
+        is_connected = self.data.last_update_success and (
-            self.data.last_update_success and self.device.state == StateType.CONNECTED
+            self.device.state == StateType.CONNECTED
+            or (not self.device.is_adopted_by_us and self.device.can_adopt)
         )
         if (
             hasattr(self, "entity_description")

homeassistant/components/unifiprotect/light.py

@@ -34,9 +34,6 @@ async def async_setup_entry(
     data: ProtectData = hass.data[DOMAIN][entry.entry_id]
 
     async def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:
-        if not device.is_adopted_by_us:
-            return
-
         if device.model == ModelType.LIGHT and device.can_write(
             data.api.bootstrap.auth_user
         ):

homeassistant/components/unifiprotect/lock.py

@@ -34,9 +34,6 @@ async def async_setup_entry(
     data: ProtectData = hass.data[DOMAIN][entry.entry_id]
 
     async def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:
-        if not device.is_adopted_by_us:
-            return
-
         if isinstance(device, Doorlock):
             async_add_entities([ProtectLock(data, device)])
 

homeassistant/components/unifiprotect/media_player.py

@@ -43,9 +43,6 @@ async def async_setup_entry(
     data: ProtectData = hass.data[DOMAIN][entry.entry_id]
 
     async def _add_new_device(device: ProtectAdoptableDeviceModel) -> None:
-        if not device.is_adopted_by_us:
-            return
-
         if isinstance(device, Camera) and device.feature_flags.has_speaker:
             async_add_entities([ProtectMediaPlayer(data, device)])
 

homeassistant/components/unifiprotect/models.py

@@ -23,6 +23,7 @@ class PermRequired(int, Enum):
 
     NO_WRITE = 1
     WRITE = 2
+    DELETE = 3
 
 
 @dataclass

homeassistant/components/unifiprotect/switch.py

@@ -320,6 +320,7 @@ async def async_setup_entry(
             data,
             ProtectPrivacyModeSwitch,
             camera_descs=[PRIVACY_MODE_SWITCH],
+            ufp_device=device,
         )
         async_add_entities(entities)
 

tests/components/unifiprotect/test_button.py

@@ -2,9 +2,9 @@
 # pylint: disable=protected-access
 from __future__ import annotations
 
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, Mock
 
-from pyunifiprotect.data.devices import Chime
+from pyunifiprotect.data.devices import Camera, Chime, Doorlock
 
 from homeassistant.components.unifiprotect.const import DEFAULT_ATTRIBUTION
 from homeassistant.const import ATTR_ATTRIBUTION, ATTR_ENTITY_ID, Platform
@@ -27,11 +27,11 @@ async def test_button_chime_remove(
     """Test removing and re-adding a light device."""
 
     await init_entry(hass, ufp, [chime])
-    assert_entity_counts(hass, Platform.BUTTON, 3, 2)
+    assert_entity_counts(hass, Platform.BUTTON, 4, 2)
     await remove_entities(hass, [chime])
     assert_entity_counts(hass, Platform.BUTTON, 0, 0)
     await adopt_devices(hass, ufp, [chime])
-    assert_entity_counts(hass, Platform.BUTTON, 3, 2)
+    assert_entity_counts(hass, Platform.BUTTON, 4, 2)
 
 
 async def test_reboot_button(
@@ -42,7 +42,7 @@ async def test_reboot_button(
     """Test button entity."""
 
     await init_entry(hass, ufp, [chime])
-    assert_entity_counts(hass, Platform.BUTTON, 3, 2)
+    assert_entity_counts(hass, Platform.BUTTON, 4, 2)
 
     ufp.api.reboot_device = AsyncMock()
 
@@ -74,7 +74,7 @@ async def test_chime_button(
     """Test button entity."""
 
     await init_entry(hass, ufp, [chime])
-    assert_entity_counts(hass, Platform.BUTTON, 3, 2)
+    assert_entity_counts(hass, Platform.BUTTON, 4, 2)
 
     ufp.api.play_speaker = AsyncMock()
 
@@ -95,3 +95,67 @@ async def test_chime_button(
         "button", "press", {ATTR_ENTITY_ID: entity_id}, blocking=True
     )
     ufp.api.play_speaker.assert_called_once()
+
+
+async def test_adopt_button(
+    hass: HomeAssistant, ufp: MockUFPFixture, doorlock: Doorlock, doorbell: Camera
+):
+    """Test button entity."""
+
+    doorlock._api = ufp.api
+    doorlock.is_adopted = False
+    doorlock.can_adopt = True
+
+    await init_entry(hass, ufp, [])
+
+    mock_msg = Mock()
+    mock_msg.changed_data = {}
+    mock_msg.old_obj = None
+    mock_msg.new_obj = doorlock
+    ufp.ws_msg(mock_msg)
+    await hass.async_block_till_done()
+
+    assert_entity_counts(hass, Platform.BUTTON, 1, 1)
+
+    ufp.api.adopt_device = AsyncMock()
+
+    unique_id = f"{doorlock.mac}_adopt"
+    entity_id = "button.test_lock_adopt_device"
+
+    entity_registry = er.async_get(hass)
+    entity = entity_registry.async_get(entity_id)
+    assert entity
+    assert not entity.disabled
+    assert entity.unique_id == unique_id
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.attributes[ATTR_ATTRIBUTION] == DEFAULT_ATTRIBUTION
+
+    await hass.services.async_call(
+        "button", "press", {ATTR_ENTITY_ID: entity_id}, blocking=True
+    )
+    ufp.api.adopt_device.assert_called_once()
+
+
+async def test_adopt_button_removed(
+    hass: HomeAssistant, ufp: MockUFPFixture, doorlock: Doorlock, doorbell: Camera
+):
+    """Test button entity."""
+
+    entity_id = "button.test_lock_adopt_device"
+    entity_registry = er.async_get(hass)
+
+    doorlock._api = ufp.api
+    doorlock.is_adopted = False
+    doorlock.can_adopt = True
+
+    await init_entry(hass, ufp, [doorlock])
+    assert_entity_counts(hass, Platform.BUTTON, 1, 1)
+    entity = entity_registry.async_get(entity_id)
+    assert entity
+
+    await adopt_devices(hass, ufp, [doorlock], fully_adopt=True)
+    assert_entity_counts(hass, Platform.BUTTON, 2, 0)
+    entity = entity_registry.async_get(entity_id)
+    assert entity is None

tests/components/unifiprotect/test_migrate.py

@@ -56,7 +56,7 @@ async def test_migrate_reboot_button(
     for entity in er.async_entries_for_config_entry(registry, ufp.entry.entry_id):
         if entity.domain == Platform.BUTTON.value:
             buttons.append(entity)
-    assert len(buttons) == 2
+    assert len(buttons) == 4
 
     assert registry.async_get(f"{Platform.BUTTON}.test_light_1_reboot_device") is None
     assert registry.async_get(f"{Platform.BUTTON}.test_light_1_reboot_device_2") is None
@@ -135,7 +135,7 @@ async def test_migrate_reboot_button_no_device(
     for entity in er.async_entries_for_config_entry(registry, ufp.entry.entry_id):
         if entity.domain == Platform.BUTTON.value:
             buttons.append(entity)
-    assert len(buttons) == 2
+    assert len(buttons) == 3
 
     entity = registry.async_get(f"{Platform.BUTTON}.unifiprotect_{light2_id.lower()}")
     assert entity is not None

tests/components/unifiprotect/utils.py

@@ -89,6 +89,7 @@ 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
 
@@ -203,10 +204,16 @@ async def adopt_devices(
     hass: HomeAssistant,
     ufp: MockUFPFixture,
     ufp_devices: list[ProtectAdoptableDeviceModel],
+    fully_adopt: bool = False,
 ):
     """Emit WS to re-adopt give Protect devices."""
 
     for ufp_device in ufp_devices:
+        if fully_adopt:
+            ufp_device.is_adopted = True
+            ufp_device.is_adopted_by_other = False
+            ufp_device.can_adopt = False
+
         mock_msg = Mock()
         mock_msg.changed_data = {}
         mock_msg.new_obj = Event(