diff --git a/homeassistant/components/unifiprotect/button.py b/homeassistant/components/unifiprotect/button.py index 9440e46b936..823a050ef09 100644 --- a/homeassistant/components/unifiprotect/button.py +++ b/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.""" diff --git a/homeassistant/components/unifiprotect/const.py b/homeassistant/components/unifiprotect/const.py index 9710279a7c4..93a0fa5ff74 100644 --- a/homeassistant/components/unifiprotect/const.py +++ b/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" diff --git a/homeassistant/components/unifiprotect/data.py b/homeassistant/components/unifiprotect/data.py index 74ef6ab37f8..d95668ea29d 100644 --- a/homeassistant/components/unifiprotect/data.py +++ b/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, @@ -156,36 +157,55 @@ class ProtectData: self._pending_camera_ids.add(camera_id) + @callback + def _async_add_device(self, device: ProtectAdoptableDeviceModel) -> None: + if device.is_adopted_by_us: + _LOGGER.debug("Device adopted: %s", device.id) + async_dispatcher_send( + 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 changed_data: + _LOGGER.debug( + "Doorbell messages updated. Updating devices with LCD screens" + ) + self.api.bootstrap.nvr.update_all_messages() + 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 or not getattr( - message.new_obj, "is_adopted_by_us", True - ): + if message.new_obj is None: 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 - ) + 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 update for all Cameras with LCD screens when NVR Doorbell settings updates - if "doorbell_settings" in message.changed_data: - _LOGGER.debug( - "Doorbell messages updated. Updating devices with LCD screens" - ) - self.api.bootstrap.nvr.update_all_messages() - for camera in self.api.bootstrap.cameras.values(): - if camera.feature_flags.has_lcd_screen: - self._async_signal_device_update(camera) # 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) - async_dispatcher_send( - self._hass, _ufpd(self._entry, DISPATCH_ADOPT), device - ) + self._async_add_device(device) elif obj.camera is not None: self._async_signal_device_update(obj.camera) elif obj.light is not None: diff --git a/homeassistant/components/unifiprotect/entity.py b/homeassistant/components/unifiprotect/entity.py index e68e5cfb81d..986dff13dc0 100644 --- a/homeassistant/components/unifiprotect/entity.py +++ b/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(data, klass, ModelType.LIGHT, light_descs) - + _async_device_entities(data, klass, ModelType.SENSOR, sense_descs) - + _async_device_entities(data, klass, ModelType.VIEWPORT, viewer_descs) - + _async_device_entities(data, klass, ModelType.DOORLOCK, lock_descs) - + _async_device_entities(data, klass, ModelType.CHIME, chime_descs) + _async_device_entities( + data, klass, ModelType.CAMERA, camera_descs, unadopted_descs + ) + + _async_device_entities( + data, klass, ModelType.LIGHT, light_descs, unadopted_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 = ( - self.data.last_update_success and self.device.state == StateType.CONNECTED + is_connected = self.data.last_update_success and ( + self.device.state == StateType.CONNECTED + or (not self.device.is_adopted_by_us and self.device.can_adopt) ) if ( hasattr(self, "entity_description") diff --git a/homeassistant/components/unifiprotect/light.py b/homeassistant/components/unifiprotect/light.py index 588b99b38d7..817e0ba1d6b 100644 --- a/homeassistant/components/unifiprotect/light.py +++ b/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 ): diff --git a/homeassistant/components/unifiprotect/lock.py b/homeassistant/components/unifiprotect/lock.py index 0a203308d1e..6a33289234e 100644 --- a/homeassistant/components/unifiprotect/lock.py +++ b/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)]) diff --git a/homeassistant/components/unifiprotect/media_player.py b/homeassistant/components/unifiprotect/media_player.py index d4046e4b8b7..f426d878ae2 100644 --- a/homeassistant/components/unifiprotect/media_player.py +++ b/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)]) diff --git a/homeassistant/components/unifiprotect/models.py b/homeassistant/components/unifiprotect/models.py index dee2006b429..adab5c032e1 100644 --- a/homeassistant/components/unifiprotect/models.py +++ b/homeassistant/components/unifiprotect/models.py @@ -23,6 +23,7 @@ class PermRequired(int, Enum): NO_WRITE = 1 WRITE = 2 + DELETE = 3 @dataclass diff --git a/homeassistant/components/unifiprotect/switch.py b/homeassistant/components/unifiprotect/switch.py index e0ddddd4c53..fa53cc3b87b 100644 --- a/homeassistant/components/unifiprotect/switch.py +++ b/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) diff --git a/tests/components/unifiprotect/test_button.py b/tests/components/unifiprotect/test_button.py index a46d74e0b8e..9db7a46dda3 100644 --- a/tests/components/unifiprotect/test_button.py +++ b/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 diff --git a/tests/components/unifiprotect/test_migrate.py b/tests/components/unifiprotect/test_migrate.py index 64c8384d400..20e4ec61ca0 100644 --- a/tests/components/unifiprotect/test_migrate.py +++ b/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 diff --git a/tests/components/unifiprotect/utils.py b/tests/components/unifiprotect/utils.py index 8d54cbddea6..0b5d29ba12a 100644 --- a/tests/components/unifiprotect/utils.py +++ b/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(