jeans/core
jeans/core
1
0
Fork 0
mirror of https://github.com/home-assistant/core.git synced 2025-07-24 13:47:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: J. Nick Koston <nick@koston.org>
This commit is contained in:
Christopher Bailey 2022-08-25 19:54:52 -04:00 committed by GitHub
parent bda9cb59d2
commit 5c0fc0c002
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
12 changed files with 246 additions and 59 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

74
homeassistant/components/unifiprotect/button.py
View File

@ -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."""

1
homeassistant/components/unifiprotect/const.py
View File

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

69
homeassistant/components/unifiprotect/data.py
View File

@ -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:

63
homeassistant/components/unifiprotect/entity.py
View File

@ -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")

3
homeassistant/components/unifiprotect/light.py
View File

@ -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
):

3
homeassistant/components/unifiprotect/lock.py
View File

@ -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)])

3
homeassistant/components/unifiprotect/media_player.py
View File

@ -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)])

1
homeassistant/components/unifiprotect/models.py
View File

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

1
homeassistant/components/unifiprotect/switch.py
View File

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

76
tests/components/unifiprotect/test_button.py
View File

@ -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

4
tests/components/unifiprotect/test_migrate.py
View File

@ -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

7
tests/components/unifiprotect/utils.py
View File

@ -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(