From b9c636ba4e240a03fbb65df6dab4d28cdfcaf78a Mon Sep 17 00:00:00 2001 From: Christopher Bailey Date: Mon, 27 Jun 2022 17:03:25 -0400 Subject: [PATCH] Automatically add newly added devices for UniFi Protect (#73879) --- .../components/unifiprotect/binary_sensor.py | 31 +++++++- .../components/unifiprotect/button.py | 18 ++++- .../components/unifiprotect/camera.py | 65 ++++++++++++----- .../components/unifiprotect/const.py | 3 + homeassistant/components/unifiprotect/data.py | 73 +++++++++++++++---- .../components/unifiprotect/entity.py | 42 ++++++++--- .../components/unifiprotect/light.py | 23 +++++- homeassistant/components/unifiprotect/lock.py | 20 ++++- .../components/unifiprotect/media_player.py | 21 ++++-- .../components/unifiprotect/number.py | 28 ++++++- .../components/unifiprotect/select.py | 22 +++++- .../components/unifiprotect/sensor.py | 37 ++++++++-- .../components/unifiprotect/switch.py | 20 ++++- .../components/unifiprotect/utils.py | 10 ++- .../unifiprotect/test_binary_sensor.py | 44 +++++++++++ tests/components/unifiprotect/test_button.py | 22 +++++- tests/components/unifiprotect/test_camera.py | 29 +++++++- tests/components/unifiprotect/test_light.py | 19 ++++- tests/components/unifiprotect/test_lock.py | 21 +++++- .../unifiprotect/test_media_player.py | 21 +++++- tests/components/unifiprotect/test_number.py | 41 +++++++++++ tests/components/unifiprotect/test_select.py | 44 +++++++++++ tests/components/unifiprotect/test_sensor.py | 30 ++++++++ tests/components/unifiprotect/test_switch.py | 30 ++++++++ tests/components/unifiprotect/utils.py | 58 ++++++++++++++- 25 files changed, 696 insertions(+), 76 deletions(-) diff --git a/homeassistant/components/unifiprotect/binary_sensor.py b/homeassistant/components/unifiprotect/binary_sensor.py index eb4b2024233..598e0632fbb 100644 --- a/homeassistant/components/unifiprotect/binary_sensor.py +++ b/homeassistant/components/unifiprotect/binary_sensor.py @@ -11,6 +11,7 @@ from pyunifiprotect.data import ( Event, Light, MountType, + ProtectAdoptableDeviceModel, ProtectModelWithId, Sensor, ) @@ -23,10 +24,11 @@ from homeassistant.components.binary_sensor import ( ) from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from .const import DISPATCH_ADOPT, DOMAIN from .data import ProtectData from .entity import ( EventThumbnailMixin, @@ -35,6 +37,7 @@ from .entity import ( async_all_device_entities, ) from .models import PermRequired, ProtectRequiredKeysMixin +from .utils import async_dispatch_id as _ufpd _LOGGER = logging.getLogger(__name__) _KEY_DOOR = "door" @@ -364,6 +367,24 @@ async def async_setup_entry( ) -> None: """Set up binary sensors for UniFi Protect integration.""" data: ProtectData = hass.data[DOMAIN][entry.entry_id] + + async def _add_new_device(device: ProtectAdoptableDeviceModel) -> None: + entities: list[ProtectDeviceEntity] = async_all_device_entities( + data, + ProtectDeviceBinarySensor, + camera_descs=CAMERA_SENSORS, + light_descs=LIGHT_SENSORS, + sense_descs=SENSE_SENSORS, + lock_descs=DOORLOCK_SENSORS, + viewer_descs=VIEWER_SENSORS, + ufp_device=device, + ) + if device.is_adopted and isinstance(device, Camera): + entities += _async_motion_entities(data, ufp_device=device) + async_add_entities(entities) + + async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device) + entities: list[ProtectDeviceEntity] = async_all_device_entities( data, ProtectDeviceBinarySensor, @@ -382,10 +403,14 @@ async def async_setup_entry( @callback def _async_motion_entities( data: ProtectData, + ufp_device: ProtectAdoptableDeviceModel | None = None, ) -> list[ProtectDeviceEntity]: entities: list[ProtectDeviceEntity] = [] - for device in data.api.bootstrap.cameras.values(): - if not device.is_adopted_by_us: + devices = ( + data.api.bootstrap.cameras.values() if ufp_device is None else [ufp_device] + ) + for device in devices: + if not device.is_adopted: continue for description in MOTION_SENSORS: diff --git a/homeassistant/components/unifiprotect/button.py b/homeassistant/components/unifiprotect/button.py index d647cdac64a..901139109d3 100644 --- a/homeassistant/components/unifiprotect/button.py +++ b/homeassistant/components/unifiprotect/button.py @@ -13,12 +13,14 @@ from homeassistant.components.button import ( ) from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant +from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from .const import 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 @dataclass @@ -79,6 +81,19 @@ async def async_setup_entry( """Discover devices on a UniFi Protect NVR.""" data: ProtectData = hass.data[DOMAIN][entry.entry_id] + async def _add_new_device(device: ProtectAdoptableDeviceModel) -> None: + entities = async_all_device_entities( + data, + ProtectButton, + all_descs=ALL_DEVICE_BUTTONS, + chime_descs=CHIME_BUTTONS, + sense_descs=SENSOR_BUTTONS, + ufp_device=device, + ) + async_add_entities(entities) + + async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device) + entities: list[ProtectDeviceEntity] = async_all_device_entities( data, ProtectButton, @@ -86,7 +101,6 @@ async def async_setup_entry( chime_descs=CHIME_BUTTONS, sense_descs=SENSOR_BUTTONS, ) - async_add_entities(entities) diff --git a/homeassistant/components/unifiprotect/camera.py b/homeassistant/components/unifiprotect/camera.py index a84346a8384..336a5ae9187 100644 --- a/homeassistant/components/unifiprotect/camera.py +++ b/homeassistant/components/unifiprotect/camera.py @@ -4,10 +4,10 @@ from __future__ import annotations from collections.abc import Generator import logging -from pyunifiprotect.api import ProtectApiClient from pyunifiprotect.data import ( Camera as UFPCamera, CameraChannel, + ProtectAdoptableDeviceModel, ProtectModelWithId, StateType, ) @@ -15,6 +15,7 @@ from pyunifiprotect.data import ( from homeassistant.components.camera import Camera, CameraEntityFeature from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity_platform import AddEntitiesCallback from .const import ( @@ -23,28 +24,39 @@ from .const import ( ATTR_FPS, ATTR_HEIGHT, ATTR_WIDTH, + DISPATCH_ADOPT, + DISPATCH_CHANNELS, DOMAIN, ) from .data import ProtectData from .entity import ProtectDeviceEntity +from .utils import async_dispatch_id as _ufpd _LOGGER = logging.getLogger(__name__) def get_camera_channels( - protect: ProtectApiClient, + data: ProtectData, + ufp_device: UFPCamera | None = None, ) -> Generator[tuple[UFPCamera, CameraChannel, bool], None, None]: """Get all the camera channels.""" - for camera in protect.bootstrap.cameras.values(): + + devices = ( + data.api.bootstrap.cameras.values() if ufp_device is None else [ufp_device] + ) + for camera in devices: if not camera.is_adopted_by_us: continue if not camera.channels: - _LOGGER.warning( - "Camera does not have any channels: %s (id: %s)", - camera.display_name, - camera.id, - ) + if ufp_device is None: + # only warn on startup + _LOGGER.warning( + "Camera does not have any channels: %s (id: %s)", + camera.display_name, + camera.id, + ) + data.async_add_pending_camera_id(camera.id) continue is_default = True @@ -60,17 +72,12 @@ def get_camera_channels( yield camera, camera.channels[0], True -async def async_setup_entry( - hass: HomeAssistant, - entry: ConfigEntry, - async_add_entities: AddEntitiesCallback, -) -> None: - """Discover cameras on a UniFi Protect NVR.""" - data: ProtectData = hass.data[DOMAIN][entry.entry_id] +def _async_camera_entities( + data: ProtectData, ufp_device: UFPCamera | None = None +) -> list[ProtectDeviceEntity]: disable_stream = data.disable_stream - - entities = [] - for camera, channel, is_default in get_camera_channels(data.api): + entities: list[ProtectDeviceEntity] = [] + for camera, channel, is_default in get_camera_channels(data, ufp_device): # do not enable streaming for package camera # 2 FPS causes a lot of buferring entities.append( @@ -95,6 +102,28 @@ async def async_setup_entry( disable_stream, ) ) + return entities + + +async def async_setup_entry( + hass: HomeAssistant, + entry: ConfigEntry, + async_add_entities: AddEntitiesCallback, +) -> None: + """Discover cameras on a UniFi Protect NVR.""" + data: ProtectData = hass.data[DOMAIN][entry.entry_id] + + async def _add_new_device(device: ProtectAdoptableDeviceModel) -> None: + if not isinstance(device, UFPCamera): + return + + entities = _async_camera_entities(data, ufp_device=device) + async_add_entities(entities) + + async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device) + async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_CHANNELS), _add_new_device) + + entities = _async_camera_entities(data) async_add_entities(entities) diff --git a/homeassistant/components/unifiprotect/const.py b/homeassistant/components/unifiprotect/const.py index 3ba22e6b85b..3c29d0c9972 100644 --- a/homeassistant/components/unifiprotect/const.py +++ b/homeassistant/components/unifiprotect/const.py @@ -60,3 +60,6 @@ PLATFORMS = [ Platform.SENSOR, Platform.SWITCH, ] + +DISPATCH_ADOPT = "adopt_device" +DISPATCH_CHANNELS = "new_camera_channels" diff --git a/homeassistant/components/unifiprotect/data.py b/homeassistant/components/unifiprotect/data.py index 4a20e816ce2..30887f04235 100644 --- a/homeassistant/components/unifiprotect/data.py +++ b/homeassistant/components/unifiprotect/data.py @@ -10,6 +10,7 @@ from pyunifiprotect import ProtectApiClient from pyunifiprotect.data import ( Bootstrap, Event, + EventType, Liveview, ModelType, ProtectAdoptableDeviceModel, @@ -20,10 +21,22 @@ from pyunifiprotect.exceptions import ClientError, NotAuthorized from homeassistant.config_entries import ConfigEntry from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback +from homeassistant.helpers.dispatcher import async_dispatcher_send from homeassistant.helpers.event import async_track_time_interval -from .const import CONF_DISABLE_RTSP, DEVICES_THAT_ADOPT, DEVICES_WITH_ENTITIES, DOMAIN -from .utils import async_get_devices, async_get_devices_by_type +from .const import ( + CONF_DISABLE_RTSP, + DEVICES_THAT_ADOPT, + DEVICES_WITH_ENTITIES, + DISPATCH_ADOPT, + DISPATCH_CHANNELS, + DOMAIN, +) +from .utils import ( + async_dispatch_id as _ufpd, + async_get_devices, + async_get_devices_by_type, +) _LOGGER = logging.getLogger(__name__) @@ -56,6 +69,7 @@ class ProtectData: self._hass = hass self._update_interval = update_interval self._subscriptions: dict[str, list[Callable[[ProtectModelWithId], None]]] = {} + self._pending_camera_ids: set[str] = set() self._unsub_interval: CALLBACK_TYPE | None = None self._unsub_websocket: CALLBACK_TYPE | None = None @@ -117,6 +131,18 @@ class ProtectData: self.last_update_success = True self._async_process_updates(updates) + @callback + def async_add_pending_camera_id(self, camera_id: str) -> None: + """ + Add pending camera. + + A "pending camera" is one that has been adopted by not had its camera channels + initialized yet. Will cause Websocket code to check for channels to be + initialized for the camera and issue a dispatch once they do. + """ + + self._pending_camera_ids.add(camera_id) + @callback def _async_process_ws_message(self, message: WSSubscriptionMessage) -> None: # removed packets are not processed yet @@ -125,8 +151,19 @@ class ProtectData: ): return - if message.new_obj.model in DEVICES_WITH_ENTITIES: - self._async_signal_device_update(message.new_obj) + obj = message.new_obj + if obj.model in DEVICES_WITH_ENTITIES: + 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 + ) + # trigger update for all Cameras with LCD screens when NVR Doorbell settings updates if "doorbell_settings" in message.changed_data: _LOGGER.debug( @@ -137,17 +174,25 @@ class ProtectData: if camera.feature_flags.has_lcd_screen: self._async_signal_device_update(camera) # trigger updates for camera that the event references - elif isinstance(message.new_obj, Event): - if message.new_obj.camera is not None: - self._async_signal_device_update(message.new_obj.camera) - elif message.new_obj.light is not None: - self._async_signal_device_update(message.new_obj.light) - elif message.new_obj.sensor is not None: - self._async_signal_device_update(message.new_obj.sensor) + elif isinstance(obj, Event): + if obj.type == EventType.DEVICE_ADOPTED: + if obj.metadata is not None and obj.metadata.device_id is not None: + device = self.api.bootstrap.get_device_from_id( + 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 + ) + elif obj.camera is not None: + self._async_signal_device_update(obj.camera) + elif obj.light is not None: + self._async_signal_device_update(obj.light) + elif obj.sensor is not None: + self._async_signal_device_update(obj.sensor) # alert user viewport needs restart so voice clients can get new options - elif len(self.api.bootstrap.viewers) > 0 and isinstance( - message.new_obj, Liveview - ): + elif len(self.api.bootstrap.viewers) > 0 and isinstance(obj, Liveview): _LOGGER.warning( "Liveviews updated. Restart Home Assistant to update Viewport select options" ) diff --git a/homeassistant/components/unifiprotect/entity.py b/homeassistant/components/unifiprotect/entity.py index 65734569de2..b7419d0a41e 100644 --- a/homeassistant/components/unifiprotect/entity.py +++ b/homeassistant/components/unifiprotect/entity.py @@ -38,12 +38,16 @@ def _async_device_entities( klass: type[ProtectDeviceEntity], model_type: ModelType, descs: Sequence[ProtectRequiredKeysMixin], + ufp_device: ProtectAdoptableDeviceModel | None = None, ) -> list[ProtectDeviceEntity]: if len(descs) == 0: return [] entities: list[ProtectDeviceEntity] = [] - for device in data.get_by_types({model_type}): + devices = ( + [ufp_device] if ufp_device is not None else data.get_by_types({model_type}) + ) + for device in devices: if not device.is_adopted_by_us: continue @@ -89,6 +93,7 @@ def async_all_device_entities( lock_descs: Sequence[ProtectRequiredKeysMixin] | None = None, chime_descs: Sequence[ProtectRequiredKeysMixin] | None = None, all_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 []) @@ -99,14 +104,33 @@ def async_all_device_entities( lock_descs = list(lock_descs or []) + all_descs chime_descs = list(chime_descs or []) + all_descs - 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) - ) + 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) + ) + + descs = [] + if ufp_device.model == ModelType.CAMERA: + descs = camera_descs + elif ufp_device.model == ModelType.LIGHT: + descs = light_descs + elif ufp_device.model == ModelType.SENSOR: + descs = sense_descs + elif ufp_device.model == ModelType.VIEWPORT: + descs = viewer_descs + elif ufp_device.model == ModelType.DOORLOCK: + descs = lock_descs + elif ufp_device.model == ModelType.CHIME: + descs = chime_descs + + if len(descs) == 0 or ufp_device.model is None: + return [] + return _async_device_entities(data, klass, ufp_device.model, descs, ufp_device) class ProtectDeviceEntity(Entity): diff --git a/homeassistant/components/unifiprotect/light.py b/homeassistant/components/unifiprotect/light.py index bd64905a289..fdfe41bca3c 100644 --- a/homeassistant/components/unifiprotect/light.py +++ b/homeassistant/components/unifiprotect/light.py @@ -4,16 +4,23 @@ from __future__ import annotations import logging from typing import Any -from pyunifiprotect.data import Light, ProtectModelWithId +from pyunifiprotect.data import ( + Light, + ModelType, + ProtectAdoptableDeviceModel, + ProtectModelWithId, +) from homeassistant.components.light import ATTR_BRIGHTNESS, ColorMode, LightEntity from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from .const import DISPATCH_ADOPT, DOMAIN from .data import ProtectData from .entity import ProtectDeviceEntity +from .utils import async_dispatch_id as _ufpd _LOGGER = logging.getLogger(__name__) @@ -25,6 +32,18 @@ async def async_setup_entry( ) -> None: """Set up lights for UniFi Protect integration.""" 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 + ): + async_add_entities([ProtectLight(data, device)]) + + async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device) + entities = [] for device in data.api.bootstrap.lights.values(): if not device.is_adopted_by_us: diff --git a/homeassistant/components/unifiprotect/lock.py b/homeassistant/components/unifiprotect/lock.py index 7258dc5f952..400d463050e 100644 --- a/homeassistant/components/unifiprotect/lock.py +++ b/homeassistant/components/unifiprotect/lock.py @@ -4,16 +4,23 @@ from __future__ import annotations import logging from typing import Any -from pyunifiprotect.data import Doorlock, LockStatusType, ProtectModelWithId +from pyunifiprotect.data import ( + Doorlock, + LockStatusType, + ProtectAdoptableDeviceModel, + ProtectModelWithId, +) from homeassistant.components.lock import LockEntity, LockEntityDescription from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from .const import DISPATCH_ADOPT, DOMAIN from .data import ProtectData from .entity import ProtectDeviceEntity +from .utils import async_dispatch_id as _ufpd _LOGGER = logging.getLogger(__name__) @@ -26,6 +33,15 @@ async def async_setup_entry( """Set up locks on a UniFi Protect NVR.""" 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)]) + + async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device) + entities = [] for device in data.api.bootstrap.doorlocks.values(): if not device.is_adopted_by_us: diff --git a/homeassistant/components/unifiprotect/media_player.py b/homeassistant/components/unifiprotect/media_player.py index b0391c9d860..41109c053f6 100644 --- a/homeassistant/components/unifiprotect/media_player.py +++ b/homeassistant/components/unifiprotect/media_player.py @@ -4,7 +4,7 @@ from __future__ import annotations import logging from typing import Any -from pyunifiprotect.data import Camera, ProtectModelWithId +from pyunifiprotect.data import Camera, ProtectAdoptableDeviceModel, ProtectModelWithId from pyunifiprotect.exceptions import StreamError from homeassistant.components import media_source @@ -23,11 +23,13 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.const import STATE_IDLE, STATE_PLAYING from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError +from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from .const import DISPATCH_ADOPT, DOMAIN from .data import ProtectData from .entity import ProtectDeviceEntity +from .utils import async_dispatch_id as _ufpd _LOGGER = logging.getLogger(__name__) @@ -40,12 +42,21 @@ async def async_setup_entry( """Discover cameras with speakers on a UniFi Protect NVR.""" 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)]) + + async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device) + entities = [] for device in data.api.bootstrap.cameras.values(): - if not device.is_adopted_by_us or not device.feature_flags.has_speaker: + if not device.is_adopted_by_us: continue - - entities.append(ProtectMediaPlayer(data, device)) + if device.feature_flags.has_speaker: + entities.append(ProtectMediaPlayer(data, device)) async_add_entities(entities) diff --git a/homeassistant/components/unifiprotect/number.py b/homeassistant/components/unifiprotect/number.py index 7bd6ce5b3d8..a017d2330b6 100644 --- a/homeassistant/components/unifiprotect/number.py +++ b/homeassistant/components/unifiprotect/number.py @@ -4,19 +4,27 @@ from __future__ import annotations from dataclasses import dataclass from datetime import timedelta -from pyunifiprotect.data import Camera, Doorlock, Light, ProtectModelWithId +from pyunifiprotect.data import ( + Camera, + Doorlock, + Light, + ProtectAdoptableDeviceModel, + ProtectModelWithId, +) from homeassistant.components.number import NumberEntity, NumberEntityDescription from homeassistant.config_entries import ConfigEntry from homeassistant.const import PERCENTAGE, TIME_SECONDS from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from .const import 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 @dataclass @@ -184,6 +192,22 @@ async def async_setup_entry( ) -> None: """Set up number entities for UniFi Protect integration.""" data: ProtectData = hass.data[DOMAIN][entry.entry_id] + + async def _add_new_device(device: ProtectAdoptableDeviceModel) -> None: + entities = async_all_device_entities( + data, + ProtectNumbers, + camera_descs=CAMERA_NUMBERS, + light_descs=LIGHT_NUMBERS, + sense_descs=SENSE_NUMBERS, + lock_descs=DOORLOCK_NUMBERS, + chime_descs=CHIME_NUMBERS, + ufp_device=device, + ) + async_add_entities(entities) + + async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device) + entities: list[ProtectDeviceEntity] = async_all_device_entities( data, ProtectNumbers, diff --git a/homeassistant/components/unifiprotect/select.py b/homeassistant/components/unifiprotect/select.py index 17bdfa390a6..5ea956ca603 100644 --- a/homeassistant/components/unifiprotect/select.py +++ b/homeassistant/components/unifiprotect/select.py @@ -19,6 +19,7 @@ from pyunifiprotect.data import ( LightModeEnableType, LightModeType, MountType, + ProtectAdoptableDeviceModel, ProtectModelWithId, RecordingMode, Sensor, @@ -32,14 +33,15 @@ from homeassistant.const import ATTR_ENTITY_ID from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import config_validation as cv, entity_platform +from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import EntityCategory from homeassistant.util.dt import utcnow -from .const import ATTR_DURATION, ATTR_MESSAGE, DOMAIN, TYPE_EMPTY_VALUE +from .const import ATTR_DURATION, ATTR_MESSAGE, DISPATCH_ADOPT, DOMAIN, TYPE_EMPTY_VALUE from .data import ProtectData from .entity import ProtectDeviceEntity, async_all_device_entities from .models import PermRequired, ProtectSetableKeysMixin, T -from .utils import async_get_light_motion_current +from .utils import async_dispatch_id as _ufpd, async_get_light_motion_current _LOGGER = logging.getLogger(__name__) _KEY_LIGHT_MOTION = "light_motion" @@ -320,6 +322,22 @@ async def async_setup_entry( ) -> None: """Set up number entities for UniFi Protect integration.""" data: ProtectData = hass.data[DOMAIN][entry.entry_id] + + async def _add_new_device(device: ProtectAdoptableDeviceModel) -> None: + entities = async_all_device_entities( + data, + ProtectSelects, + camera_descs=CAMERA_SELECTS, + light_descs=LIGHT_SELECTS, + sense_descs=SENSE_SELECTS, + viewer_descs=VIEWER_SELECTS, + lock_descs=DOORLOCK_SELECTS, + ufp_device=device, + ) + async_add_entities(entities) + + async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device) + entities: list[ProtectDeviceEntity] = async_all_device_entities( data, ProtectSelects, diff --git a/homeassistant/components/unifiprotect/sensor.py b/homeassistant/components/unifiprotect/sensor.py index 012d52ae215..a46e4c790b7 100644 --- a/homeassistant/components/unifiprotect/sensor.py +++ b/homeassistant/components/unifiprotect/sensor.py @@ -36,10 +36,11 @@ from homeassistant.const import ( TIME_SECONDS, ) from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from .const import DISPATCH_ADOPT, DOMAIN from .data import ProtectData from .entity import ( EventThumbnailMixin, @@ -48,7 +49,7 @@ from .entity import ( async_all_device_entities, ) from .models import PermRequired, ProtectRequiredKeysMixin, T -from .utils import async_get_light_motion_current +from .utils import async_dispatch_id as _ufpd, async_get_light_motion_current _LOGGER = logging.getLogger(__name__) OBJECT_TYPE_NONE = "none" @@ -594,6 +595,26 @@ async def async_setup_entry( ) -> None: """Set up sensors for UniFi Protect integration.""" data: ProtectData = hass.data[DOMAIN][entry.entry_id] + + async def _add_new_device(device: ProtectAdoptableDeviceModel) -> None: + entities = async_all_device_entities( + data, + ProtectDeviceSensor, + all_descs=ALL_DEVICES_SENSORS, + camera_descs=CAMERA_SENSORS + CAMERA_DISABLED_SENSORS, + sense_descs=SENSE_SENSORS, + light_descs=LIGHT_SENSORS, + lock_descs=DOORLOCK_SENSORS, + chime_descs=CHIME_SENSORS, + viewer_descs=VIEWER_SENSORS, + ufp_device=device, + ) + if device.is_adopted_by_us and isinstance(device, Camera): + entities += _async_motion_entities(data, ufp_device=device) + async_add_entities(entities) + + async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device) + entities: list[ProtectDeviceEntity] = async_all_device_entities( data, ProtectDeviceSensor, @@ -614,13 +635,17 @@ async def async_setup_entry( @callback def _async_motion_entities( data: ProtectData, + ufp_device: Camera | None = None, ) -> list[ProtectDeviceEntity]: entities: list[ProtectDeviceEntity] = [] - for device in data.api.bootstrap.cameras.values(): - for description in MOTION_TRIP_SENSORS: - if not device.is_adopted_by_us: - continue + devices = ( + data.api.bootstrap.cameras.values() if ufp_device is None else [ufp_device] + ) + for device in devices: + if not device.is_adopted_by_us: + continue + for description in MOTION_TRIP_SENSORS: entities.append(ProtectDeviceSensor(data, device, description)) _LOGGER.debug( "Adding trip sensor entity %s for %s", diff --git a/homeassistant/components/unifiprotect/switch.py b/homeassistant/components/unifiprotect/switch.py index 8b3661ce324..71812459b95 100644 --- a/homeassistant/components/unifiprotect/switch.py +++ b/homeassistant/components/unifiprotect/switch.py @@ -15,13 +15,15 @@ from pyunifiprotect.data import ( from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant +from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from .const import 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__) @@ -302,6 +304,22 @@ async def async_setup_entry( ) -> None: """Set up sensors for UniFi Protect integration.""" data: ProtectData = hass.data[DOMAIN][entry.entry_id] + + async def _add_new_device(device: ProtectAdoptableDeviceModel) -> None: + entities = async_all_device_entities( + data, + ProtectSwitch, + camera_descs=CAMERA_SWITCHES, + light_descs=LIGHT_SWITCHES, + sense_descs=SENSE_SWITCHES, + lock_descs=DOORLOCK_SWITCHES, + viewer_descs=VIEWER_SWITCHES, + ufp_device=device, + ) + async_add_entities(entities) + + async_dispatcher_connect(hass, _ufpd(entry, DISPATCH_ADOPT), _add_new_device) + entities: list[ProtectDeviceEntity] = async_all_device_entities( data, ProtectSwitch, diff --git a/homeassistant/components/unifiprotect/utils.py b/homeassistant/components/unifiprotect/utils.py index 72baab334f3..808117aac9e 100644 --- a/homeassistant/components/unifiprotect/utils.py +++ b/homeassistant/components/unifiprotect/utils.py @@ -15,9 +15,10 @@ from pyunifiprotect.data import ( ProtectAdoptableDeviceModel, ) +from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback -from .const import ModelType +from .const import DOMAIN, ModelType def get_nested_attr(obj: Any, attr: str) -> Any: @@ -98,3 +99,10 @@ def async_get_light_motion_current(obj: Light) -> str: ): return f"{LightModeType.MOTION.value}Dark" return obj.light_mode_settings.mode.value + + +@callback +def async_dispatch_id(entry: ConfigEntry, dispatch: str) -> str: + """Generate entry specific dispatch ID.""" + + return f"{DOMAIN}.{entry.entry_id}.{dispatch}" diff --git a/tests/components/unifiprotect/test_binary_sensor.py b/tests/components/unifiprotect/test_binary_sensor.py index 856c034905f..640bf81ec49 100644 --- a/tests/components/unifiprotect/test_binary_sensor.py +++ b/tests/components/unifiprotect/test_binary_sensor.py @@ -32,15 +32,59 @@ from homeassistant.helpers import entity_registry as er from .utils import ( MockUFPFixture, + adopt_devices, assert_entity_counts, ids_from_device_description, init_entry, + remove_entities, ) LIGHT_SENSOR_WRITE = LIGHT_SENSORS[:2] SENSE_SENSORS_WRITE = SENSE_SENSORS[:4] +async def test_binary_sensor_camera_remove( + hass: HomeAssistant, ufp: MockUFPFixture, doorbell: Camera, unadopted_camera: Camera +): + """Test removing and re-adding a camera device.""" + + 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]) + 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) + + +async def test_binary_sensor_light_remove( + hass: HomeAssistant, ufp: MockUFPFixture, light: Light +): + """Test removing and re-adding a light device.""" + + 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]) + assert_entity_counts(hass, Platform.BINARY_SENSOR, 0, 0) + await adopt_devices(hass, ufp, [light]) + assert_entity_counts(hass, Platform.BINARY_SENSOR, 2, 2) + + +async def test_binary_sensor_sensor_remove( + hass: HomeAssistant, ufp: MockUFPFixture, sensor_all: Sensor +): + """Test removing and re-adding a light device.""" + + 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]) + 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) + + async def test_binary_sensor_setup_light( hass: HomeAssistant, ufp: MockUFPFixture, light: Light ): diff --git a/tests/components/unifiprotect/test_button.py b/tests/components/unifiprotect/test_button.py index a846214a7aa..a46d74e0b8e 100644 --- a/tests/components/unifiprotect/test_button.py +++ b/tests/components/unifiprotect/test_button.py @@ -11,7 +11,27 @@ from homeassistant.const import ATTR_ATTRIBUTION, ATTR_ENTITY_ID, Platform from homeassistant.core import HomeAssistant from homeassistant.helpers import entity_registry as er -from .utils import MockUFPFixture, assert_entity_counts, enable_entity, init_entry +from .utils import ( + MockUFPFixture, + adopt_devices, + assert_entity_counts, + enable_entity, + init_entry, + remove_entities, +) + + +async def test_button_chime_remove( + hass: HomeAssistant, ufp: MockUFPFixture, chime: Chime +): + """Test removing and re-adding a light device.""" + + await init_entry(hass, ufp, [chime]) + assert_entity_counts(hass, Platform.BUTTON, 3, 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) async def test_reboot_button( diff --git a/tests/components/unifiprotect/test_camera.py b/tests/components/unifiprotect/test_camera.py index 6fad7cb899e..2b103e8d714 100644 --- a/tests/components/unifiprotect/test_camera.py +++ b/tests/components/unifiprotect/test_camera.py @@ -33,9 +33,11 @@ from homeassistant.setup import async_setup_component from .utils import ( MockUFPFixture, + adopt_devices, assert_entity_counts, enable_entity, init_entry, + remove_entities, time_changed, ) @@ -268,18 +270,37 @@ async def test_basic_setup( await validate_no_stream_camera_state(hass, doorbell, 3, entity_id, features=0) -async def test_missing_channels( - hass: HomeAssistant, ufp: MockUFPFixture, camera: ProtectCamera -): +async def test_adopt(hass: HomeAssistant, ufp: MockUFPFixture, camera: ProtectCamera): """Test setting up camera with no camera channels.""" camera1 = camera.copy() camera1.channels = [] await init_entry(hass, ufp, [camera1]) - assert_entity_counts(hass, Platform.CAMERA, 0, 0) + await remove_entities(hass, [camera1]) + assert_entity_counts(hass, Platform.CAMERA, 0, 0) + camera1.channels = [] + await adopt_devices(hass, ufp, [camera1]) + assert_entity_counts(hass, Platform.CAMERA, 0, 0) + + camera1.channels = camera.channels + for channel in camera1.channels: + channel._api = ufp.api + + mock_msg = Mock() + mock_msg.changed_data = {"channels": camera.channels} + mock_msg.new_obj = camera1 + ufp.ws_msg(mock_msg) + await hass.async_block_till_done() + assert_entity_counts(hass, Platform.CAMERA, 2, 1) + + await remove_entities(hass, [camera1]) + assert_entity_counts(hass, Platform.CAMERA, 0, 0) + await adopt_devices(hass, ufp, [camera1]) + assert_entity_counts(hass, Platform.CAMERA, 2, 1) + async def test_camera_image( hass: HomeAssistant, ufp: MockUFPFixture, camera: ProtectCamera diff --git a/tests/components/unifiprotect/test_light.py b/tests/components/unifiprotect/test_light.py index 3c575de8d00..40f2191828e 100644 --- a/tests/components/unifiprotect/test_light.py +++ b/tests/components/unifiprotect/test_light.py @@ -19,7 +19,24 @@ from homeassistant.const import ( from homeassistant.core import HomeAssistant from homeassistant.helpers import entity_registry as er -from .utils import MockUFPFixture, assert_entity_counts, init_entry +from .utils import ( + MockUFPFixture, + adopt_devices, + assert_entity_counts, + init_entry, + remove_entities, +) + + +async def test_light_remove(hass: HomeAssistant, ufp: MockUFPFixture, light: Light): + """Test removing and re-adding a light device.""" + + await init_entry(hass, ufp, [light]) + assert_entity_counts(hass, Platform.LIGHT, 1, 1) + await remove_entities(hass, [light]) + assert_entity_counts(hass, Platform.LIGHT, 0, 0) + await adopt_devices(hass, ufp, [light]) + assert_entity_counts(hass, Platform.LIGHT, 1, 1) async def test_light_setup( diff --git a/tests/components/unifiprotect/test_lock.py b/tests/components/unifiprotect/test_lock.py index 21b3c77deb5..d6534e93845 100644 --- a/tests/components/unifiprotect/test_lock.py +++ b/tests/components/unifiprotect/test_lock.py @@ -21,7 +21,26 @@ from homeassistant.const import ( from homeassistant.core import HomeAssistant from homeassistant.helpers import entity_registry as er -from .utils import MockUFPFixture, assert_entity_counts, init_entry +from .utils import ( + MockUFPFixture, + adopt_devices, + assert_entity_counts, + init_entry, + remove_entities, +) + + +async def test_lock_remove( + hass: HomeAssistant, ufp: MockUFPFixture, doorlock: Doorlock +): + """Test removing and re-adding a lock device.""" + + await init_entry(hass, ufp, [doorlock]) + assert_entity_counts(hass, Platform.LOCK, 1, 1) + await remove_entities(hass, [doorlock]) + assert_entity_counts(hass, Platform.LOCK, 0, 0) + await adopt_devices(hass, ufp, [doorlock]) + assert_entity_counts(hass, Platform.LOCK, 1, 1) async def test_lock_setup( diff --git a/tests/components/unifiprotect/test_media_player.py b/tests/components/unifiprotect/test_media_player.py index 678fa0c9be4..ade84e2d51c 100644 --- a/tests/components/unifiprotect/test_media_player.py +++ b/tests/components/unifiprotect/test_media_player.py @@ -25,7 +25,26 @@ from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import entity_registry as er -from .utils import MockUFPFixture, assert_entity_counts, init_entry +from .utils import ( + MockUFPFixture, + adopt_devices, + assert_entity_counts, + init_entry, + remove_entities, +) + + +async def test_media_player_camera_remove( + hass: HomeAssistant, ufp: MockUFPFixture, doorbell: Camera +): + """Test removing and re-adding a light device.""" + + await init_entry(hass, ufp, [doorbell]) + assert_entity_counts(hass, Platform.MEDIA_PLAYER, 1, 1) + await remove_entities(hass, [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) async def test_media_player_setup( diff --git a/tests/components/unifiprotect/test_number.py b/tests/components/unifiprotect/test_number.py index 656f7d08ba5..51e9dfc85a2 100644 --- a/tests/components/unifiprotect/test_number.py +++ b/tests/components/unifiprotect/test_number.py @@ -21,12 +21,53 @@ from homeassistant.helpers import entity_registry as er from .utils import ( MockUFPFixture, + adopt_devices, assert_entity_counts, ids_from_device_description, init_entry, + remove_entities, ) +async def test_number_sensor_camera_remove( + hass: HomeAssistant, ufp: MockUFPFixture, camera: Camera, unadopted_camera: Camera +): + """Test removing and re-adding a camera device.""" + + await init_entry(hass, ufp, [camera, unadopted_camera]) + assert_entity_counts(hass, Platform.NUMBER, 3, 3) + await remove_entities(hass, [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) + + +async def test_number_sensor_light_remove( + hass: HomeAssistant, ufp: MockUFPFixture, light: Light +): + """Test removing and re-adding a light device.""" + + await init_entry(hass, ufp, [light]) + assert_entity_counts(hass, Platform.NUMBER, 2, 2) + await remove_entities(hass, [light]) + assert_entity_counts(hass, Platform.NUMBER, 0, 0) + await adopt_devices(hass, ufp, [light]) + assert_entity_counts(hass, Platform.NUMBER, 2, 2) + + +async def test_number_lock_remove( + hass: HomeAssistant, ufp: MockUFPFixture, doorlock: Doorlock +): + """Test removing and re-adding a light device.""" + + await init_entry(hass, ufp, [doorlock]) + assert_entity_counts(hass, Platform.NUMBER, 1, 1) + await remove_entities(hass, [doorlock]) + assert_entity_counts(hass, Platform.NUMBER, 0, 0) + await adopt_devices(hass, ufp, [doorlock]) + assert_entity_counts(hass, Platform.NUMBER, 1, 1) + + async def test_number_setup_light( hass: HomeAssistant, ufp: MockUFPFixture, light: Light ): diff --git a/tests/components/unifiprotect/test_select.py b/tests/components/unifiprotect/test_select.py index 637a0d4ad5d..46bc70f61f6 100644 --- a/tests/components/unifiprotect/test_select.py +++ b/tests/components/unifiprotect/test_select.py @@ -41,12 +41,56 @@ from homeassistant.helpers import entity_registry as er from .utils import ( MockUFPFixture, + adopt_devices, assert_entity_counts, ids_from_device_description, init_entry, + remove_entities, ) +async def test_select_camera_remove( + hass: HomeAssistant, ufp: MockUFPFixture, doorbell: Camera, unadopted_camera: Camera +): + """Test removing and re-adding a camera device.""" + + 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]) + assert_entity_counts(hass, Platform.SELECT, 0, 0) + await adopt_devices(hass, ufp, [doorbell, unadopted_camera]) + assert_entity_counts(hass, Platform.SELECT, 4, 4) + + +async def test_select_light_remove( + hass: HomeAssistant, ufp: MockUFPFixture, light: Light +): + """Test removing and re-adding a light device.""" + + 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]) + assert_entity_counts(hass, Platform.SELECT, 0, 0) + await adopt_devices(hass, ufp, [light]) + assert_entity_counts(hass, Platform.SELECT, 2, 2) + + +async def test_select_viewer_remove( + hass: HomeAssistant, ufp: MockUFPFixture, viewer: Viewer +): + """Test removing and re-adding a light device.""" + + 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]) + assert_entity_counts(hass, Platform.SELECT, 0, 0) + await adopt_devices(hass, ufp, [viewer]) + assert_entity_counts(hass, Platform.SELECT, 1, 1) + + async def test_select_setup_light( hass: HomeAssistant, ufp: MockUFPFixture, light: Light ): diff --git a/tests/components/unifiprotect/test_sensor.py b/tests/components/unifiprotect/test_sensor.py index e204b09b1b0..fcad6ce2725 100644 --- a/tests/components/unifiprotect/test_sensor.py +++ b/tests/components/unifiprotect/test_sensor.py @@ -41,10 +41,12 @@ from homeassistant.helpers import entity_registry as er from .utils import ( MockUFPFixture, + adopt_devices, assert_entity_counts, enable_entity, ids_from_device_description, init_entry, + remove_entities, reset_objects, time_changed, ) @@ -53,6 +55,34 @@ CAMERA_SENSORS_WRITE = CAMERA_SENSORS[:5] SENSE_SENSORS_WRITE = SENSE_SENSORS[:8] +async def test_sensor_camera_remove( + hass: HomeAssistant, ufp: MockUFPFixture, doorbell: Camera, unadopted_camera: Camera +): + """Test removing and re-adding a camera device.""" + + 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]) + assert_entity_counts(hass, Platform.SENSOR, 12, 9) + await adopt_devices(hass, ufp, [doorbell, unadopted_camera]) + assert_entity_counts(hass, Platform.SENSOR, 25, 13) + + +async def test_sensor_sensor_remove( + hass: HomeAssistant, ufp: MockUFPFixture, sensor_all: Sensor +): + """Test removing and re-adding a light device.""" + + 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]) + assert_entity_counts(hass, Platform.SENSOR, 12, 9) + await adopt_devices(hass, ufp, [sensor_all]) + assert_entity_counts(hass, Platform.SENSOR, 22, 14) + + async def test_sensor_setup_sensor( hass: HomeAssistant, ufp: MockUFPFixture, sensor_all: Sensor ): diff --git a/tests/components/unifiprotect/test_switch.py b/tests/components/unifiprotect/test_switch.py index 0c45ec28b7b..684e3b8e441 100644 --- a/tests/components/unifiprotect/test_switch.py +++ b/tests/components/unifiprotect/test_switch.py @@ -19,10 +19,12 @@ from homeassistant.helpers import entity_registry as er from .utils import ( MockUFPFixture, + adopt_devices, assert_entity_counts, enable_entity, ids_from_device_description, init_entry, + remove_entities, ) CAMERA_SWITCHES_BASIC = [ @@ -37,6 +39,34 @@ CAMERA_SWITCHES_NO_EXTRA = [ ] +async def test_switch_camera_remove( + hass: HomeAssistant, ufp: MockUFPFixture, doorbell: Camera, unadopted_camera: Camera +): + """Test removing and re-adding a camera device.""" + + 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]) + assert_entity_counts(hass, Platform.SWITCH, 0, 0) + await adopt_devices(hass, ufp, [doorbell, unadopted_camera]) + assert_entity_counts(hass, Platform.SWITCH, 13, 12) + + +async def test_switch_light_remove( + hass: HomeAssistant, ufp: MockUFPFixture, light: Light +): + """Test removing and re-adding a light device.""" + + 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]) + assert_entity_counts(hass, Platform.SWITCH, 0, 0) + await adopt_devices(hass, ufp, [light]) + assert_entity_counts(hass, Platform.SWITCH, 2, 1) + + async def test_switch_setup_no_perm( hass: HomeAssistant, ufp: MockUFPFixture, diff --git a/tests/components/unifiprotect/utils.py b/tests/components/unifiprotect/utils.py index 517da9e73c6..260c6996128 100644 --- a/tests/components/unifiprotect/utils.py +++ b/tests/components/unifiprotect/utils.py @@ -5,11 +5,15 @@ from __future__ import annotations from dataclasses import dataclass from datetime import timedelta from typing import Any, Callable, Sequence +from unittest.mock import Mock from pyunifiprotect import ProtectApiClient from pyunifiprotect.data import ( Bootstrap, Camera, + Event, + EventType, + ModelType, ProtectAdoptableDeviceModel, WSSubscriptionMessage, ) @@ -18,7 +22,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 entity_registry as er +from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.helpers.entity import EntityDescription import homeassistant.util.dt as dt_util @@ -166,3 +170,55 @@ async def init_entry( await hass.config_entries.async_setup(ufp.entry.entry_id) await hass.async_block_till_done() + + +async def remove_entities( + hass: HomeAssistant, + 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() + entity = entity_registry.async_get(f"{Platform.SENSOR}.{name}_uptime") + assert entity is not None + + device_id = entity.device_id + for reg in list(entity_registry.entities.values()): + if reg.device_id == device_id: + entity_registry.async_remove(reg.entity_id) + device_registry.async_remove_device(device_id) + + await hass.async_block_till_done() + + +async def adopt_devices( + hass: HomeAssistant, + ufp: MockUFPFixture, + ufp_devices: list[ProtectAdoptableDeviceModel], +): + """Emit WS to re-adopt give Protect devices.""" + + for ufp_device in ufp_devices: + mock_msg = Mock() + mock_msg.changed_data = {} + mock_msg.new_obj = Event( + api=ufp_device.api, + id=random_hex(24), + smart_detect_types=[], + smart_detect_event_ids=[], + type=EventType.DEVICE_ADOPTED, + start=dt_util.utcnow(), + score=100, + metadata={"device_id": ufp_device.id}, + model=ModelType.EVENT, + ) + ufp.ws_msg(mock_msg) + + await hass.async_block_till_done()