Add UniFi Protect global services (#63768)

Co-authored-by: J. Nick Koston <nick@koston.org>
This commit is contained in:
Christopher Bailey 2022-01-10 16:04:53 -05:00 committed by GitHub
parent 52959cf48c
commit 0030f114f9
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 464 additions and 14 deletions

View file

@ -3,13 +3,14 @@ from __future__ import annotations
import asyncio
from datetime import timedelta
import functools
import logging
from aiohttp import CookieJar
from aiohttp.client_exceptions import ServerDisconnectedError
from pyunifiprotect import NotAuthorized, NvrError, ProtectApiClient
from homeassistant.config_entries import ConfigEntry
from homeassistant.config_entries import ConfigEntry, ConfigEntryState
from homeassistant.const import (
CONF_HOST,
CONF_PASSWORD,
@ -23,16 +24,22 @@ from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
from homeassistant.helpers.aiohttp_client import async_create_clientsession
from .const import (
ALL_GLOBAL_SERIVCES,
CONF_ALL_UPDATES,
CONF_OVERRIDE_CHOST,
DEFAULT_SCAN_INTERVAL,
DEVICES_FOR_SUBSCRIBE,
DOMAIN,
DOORBELL_TEXT_SCHEMA,
MIN_REQUIRED_PROTECT_V,
OUTDATED_LOG_MESSAGE,
PLATFORMS,
SERVICE_ADD_DOORBELL_TEXT,
SERVICE_REMOVE_DOORBELL_TEXT,
SERVICE_SET_DEFAULT_DOORBELL_TEXT,
)
from .data import ProtectData
from .services import add_doorbell_text, remove_doorbell_text, set_default_doorbell_text
from .views import ThumbnailProxyView
_LOGGER = logging.getLogger(__name__)
@ -83,6 +90,28 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
hass.data.setdefault(DOMAIN, {})[entry.entry_id] = data_service
hass.config_entries.async_setup_platforms(entry, PLATFORMS)
services = [
(
SERVICE_ADD_DOORBELL_TEXT,
functools.partial(add_doorbell_text, hass),
DOORBELL_TEXT_SCHEMA,
),
(
SERVICE_REMOVE_DOORBELL_TEXT,
functools.partial(remove_doorbell_text, hass),
DOORBELL_TEXT_SCHEMA,
),
(
SERVICE_SET_DEFAULT_DOORBELL_TEXT,
functools.partial(set_default_doorbell_text, hass),
DOORBELL_TEXT_SCHEMA,
),
]
for name, method, schema in services:
if hass.services.has_service(DOMAIN, name):
continue
hass.services.async_register(DOMAIN, name, method, schema=schema)
hass.http.register_view(ThumbnailProxyView(hass))
entry.async_on_unload(entry.add_update_listener(_async_options_updated))
@ -104,4 +133,14 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
data: ProtectData = hass.data[DOMAIN][entry.entry_id]
await data.async_stop()
hass.data[DOMAIN].pop(entry.entry_id)
return unload_ok
loaded_entries = [
entry
for entry in hass.config_entries.async_entries(DOMAIN)
if entry.state == ConfigEntryState.LOADED
]
if len(loaded_entries) == 1:
for name in ALL_GLOBAL_SERIVCES:
hass.services.async_remove(DOMAIN, name)
return bool(unload_ok)

View file

@ -1,8 +1,10 @@
"""Constant definitions for UniFi Protect Integration."""
from pyunifiprotect.data.types import ModelType, Version
import voluptuous as vol
from homeassistant.const import Platform
from homeassistant.const import ATTR_DEVICE_ID, ATTR_ENTITY_ID, Platform
from homeassistant.helpers import config_validation as cv
DOMAIN = "unifiprotect"
@ -15,6 +17,7 @@ ATTR_BITRATE = "bitrate"
ATTR_CHANNEL_ID = "channel_id"
ATTR_MESSAGE = "message"
ATTR_DURATION = "duration"
ATTR_ANONYMIZE = "anonymize"
CONF_DISABLE_RTSP = "disable_rtsp"
CONF_ALL_UPDATES = "all_updates"
@ -46,6 +49,16 @@ OUTDATED_LOG_MESSAGE = "You are running v%s of UniFi Protect. Minimum required v
TYPE_EMPTY_VALUE = ""
SERVICE_ADD_DOORBELL_TEXT = "add_doorbell_text"
SERVICE_REMOVE_DOORBELL_TEXT = "remove_doorbell_text"
SERVICE_SET_DEFAULT_DOORBELL_TEXT = "set_default_doorbell_text"
ALL_GLOBAL_SERIVCES = [
SERVICE_ADD_DOORBELL_TEXT,
SERVICE_REMOVE_DOORBELL_TEXT,
SERVICE_SET_DEFAULT_DOORBELL_TEXT,
]
PLATFORMS = [
Platform.BINARY_SENSOR,
Platform.BUTTON,
@ -57,3 +70,43 @@ PLATFORMS = [
Platform.SENSOR,
Platform.SWITCH,
]
DOORBELL_TEXT_SCHEMA = vol.All(
vol.Schema(
{
**cv.ENTITY_SERVICE_FIELDS,
vol.Required(ATTR_MESSAGE): cv.string,
},
),
cv.has_at_least_one_key(ATTR_DEVICE_ID),
)
GENERATE_DATA_SCHEMA = vol.All(
vol.Schema(
{
**cv.ENTITY_SERVICE_FIELDS,
vol.Required(ATTR_DURATION): vol.Coerce(int),
vol.Required(ATTR_ANONYMIZE): vol.Coerce(bool),
},
),
cv.has_at_least_one_key(ATTR_DEVICE_ID),
)
PROFILE_WS_SCHEMA = vol.All(
vol.Schema(
{
**cv.ENTITY_SERVICE_FIELDS,
vol.Required(ATTR_DURATION): vol.Coerce(int),
},
),
cv.has_at_least_one_key(ATTR_DEVICE_ID),
)
SET_DOORBELL_LCD_MESSAGE_SCHEMA = vol.Schema(
{
vol.Required(ATTR_ENTITY_ID): cv.entity_ids,
vol.Required(ATTR_MESSAGE): cv.string,
vol.Optional(ATTR_DURATION, default="None"): cv.string,
}
)

View file

@ -0,0 +1,112 @@
"""UniFi Protect Integration services."""
from __future__ import annotations
import asyncio
from typing import Any
from pydantic import ValidationError
from pyunifiprotect.api import ProtectApiClient
from pyunifiprotect.exceptions import BadRequest
from homeassistant.core import HomeAssistant, ServiceCall, callback
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers import device_registry as dr
from homeassistant.helpers.service import async_extract_referenced_entity_ids
from .const import ATTR_MESSAGE, DOMAIN
from .data import ProtectData
def _async_all_ufp_instances(hass: HomeAssistant) -> list[ProtectApiClient]:
"""All active UFP instances."""
return [
data.api for data in hass.data[DOMAIN].values() if isinstance(data, ProtectData)
]
@callback
def _async_unifi_mac_from_hass(mac: str) -> str:
# MAC addresses in UFP are always caps
return mac.replace(":", "").upper()
@callback
def _async_get_macs_for_device(device_entry: dr.DeviceEntry) -> list[str]:
return [
_async_unifi_mac_from_hass(cval)
for ctype, cval in device_entry.connections
if ctype == dr.CONNECTION_NETWORK_MAC
]
@callback
def _async_get_ufp_instances(
hass: HomeAssistant, device_id: str
) -> tuple[dr.DeviceEntry, ProtectApiClient]:
device_registry = dr.async_get(hass)
if not (device_entry := device_registry.async_get(device_id)):
raise HomeAssistantError(f"No device found for device id: {device_id}")
if device_entry.via_device_id is not None:
return _async_get_ufp_instances(hass, device_entry.via_device_id)
macs = _async_get_macs_for_device(device_entry)
ufp_instances = [
i for i in _async_all_ufp_instances(hass) if i.bootstrap.nvr.mac in macs
]
if not ufp_instances:
# should not be possible unless user manually enters a bad device ID
raise HomeAssistantError( # pragma: no cover
f"No UniFi Protect NVR found for device ID: {device_id}"
)
return device_entry, ufp_instances[0]
@callback
def _async_get_protect_from_call(
hass: HomeAssistant, call: ServiceCall
) -> list[tuple[dr.DeviceEntry, ProtectApiClient]]:
referenced = async_extract_referenced_entity_ids(hass, call)
instances: list[tuple[dr.DeviceEntry, ProtectApiClient]] = []
for device_id in referenced.referenced_devices:
instances.append(_async_get_ufp_instances(hass, device_id))
return instances
async def _async_call_nvr(
instances: list[tuple[dr.DeviceEntry, ProtectApiClient]],
method: str,
*args: Any,
**kwargs: Any,
) -> None:
try:
await asyncio.gather(
*(getattr(i.bootstrap.nvr, method)(*args, **kwargs) for _, i in instances)
)
except (BadRequest, ValidationError) as err:
raise HomeAssistantError(str(err)) from err
async def add_doorbell_text(hass: HomeAssistant, call: ServiceCall) -> None:
"""Add a custom doorbell text message."""
message: str = call.data[ATTR_MESSAGE]
instances = _async_get_protect_from_call(hass, call)
await _async_call_nvr(instances, "add_custom_doorbell_message", message)
async def remove_doorbell_text(hass: HomeAssistant, call: ServiceCall) -> None:
"""Remove a custom doorbell text message."""
message: str = call.data[ATTR_MESSAGE]
instances = _async_get_protect_from_call(hass, call)
await _async_call_nvr(instances, "remove_custom_doorbell_message", message)
async def set_default_doorbell_text(hass: HomeAssistant, call: ServiceCall) -> None:
"""Set the default doorbell text message."""
message: str = call.data[ATTR_MESSAGE]
instances = _async_get_protect_from_call(hass, call)
await _async_call_nvr(instances, "set_default_doorbell_message", message)

View file

@ -1,30 +1,81 @@
add_doorbell_text:
name: Add Custom Doorbell Text
description: Adds a new custom message for Doorbells.
fields:
device_id:
name: UniFi Protect NVR
description: Any device from the UniFi Protect instance you want to change. In case you have multiple Protect Instances.
required: true
selector:
device:
integration: unifiprotect
message:
name: Custom Message
description: New custom message to add for Doorbells. Must be less than 30 characters.
example: Come In
required: true
selector:
text:
remove_doorbell_text:
name: Remove Custom Doorbell Text
description: Removes an existing message for Doorbells.
fields:
device_id:
name: UniFi Protect NVR
description: Any device from the UniFi Protect instance you want to change. In case you have multiple Protect Instances.
required: true
selector:
device:
integration: unifiprotect
message:
name: Custom Message
description: Existing custom message to remove for Doorbells.
example: Go Away!
required: true
selector:
text:
set_default_doorbell_text:
name: Set Default Doorbell Text
description: Sets the default doorbell message. This will be the message that is automatically selected when a message "expires".
fields:
device_id:
name: UniFi Protect NVR
description: Any device from the UniFi Protect instance you want to change. In case you have multiple Protect Instances.
required: true
selector:
device:
integration: unifiprotect
message:
name: Default Message
description: The default message for your Doorbell. Must be less than 30 characters.
example: Welcome!
required: true
selector:
text:
set_doorbell_message:
name: Set Doorbell message
description: >
Use to dynamically set the message on a Doorbell LCD screen. Should only be used to set dynamic messages
(i.e. setting the current outdoor temperature on your Doorbell). Static messages should still using the Select entity and the
add_doorbell_text / remove_doorbell_text services.
Use to dynamically set the message on a Doorbell LCD screen. This service should only be used to set dynamic messages (i.e. setting the current outdoor temperature on your Doorbell). Static messages should still be set using the Select entity and can be added/removed using the add_doorbell_text/remove_doorbell_text services.
fields:
entity_id:
name: Doorbell Text
description: (Required) Doorbell to display message on
description: The Doorbell Text select entity for your Doorbell.
example: "select.front_doorbell_camera_doorbell_text"
required: true
selector:
entity:
integration: unifiprotect
domain: select
device_class: unifiprotect__lcd_message
message:
name: Message to display
description: (Required) Message to display on LCD Panel. Max 30 characters
description: The message you would like to display on the LCD screen of your Doorbell. Must be less than 30 characters.
example: "Welcome | 09:23 | 25°C"
required: true
selector:
text:
duration:
name: Duration (minutes)
description: "(Optional) Number of minutes to display message, before returning to default. Leave blank to display always"
name: Duration
description: Number of minutes to display the message for before returning to the default message. The default is to not expire.
example: 5
selector:
number:
@ -32,3 +83,4 @@ set_doorbell_message:
max: 120
step: 1
mode: slider
unit_of_measurement: minutes

View file

@ -1,16 +1,18 @@
"""Test the UniFi Protect setup flow."""
from __future__ import annotations
from unittest.mock import AsyncMock
from unittest.mock import AsyncMock, patch
from pyunifiprotect import NotAuthorized, NvrError
from pyunifiprotect.data import NVR
from homeassistant.components.unifiprotect.const import CONF_DISABLE_RTSP
from homeassistant.components.unifiprotect.const import CONF_DISABLE_RTSP, DOMAIN
from homeassistant.config_entries import ConfigEntryState
from homeassistant.core import HomeAssistant
from .conftest import MockEntityFixture
from .conftest import MockBootstrap, MockEntityFixture
from tests.common import MockConfigEntry
async def test_setup(hass: HomeAssistant, mock_entry: MockEntityFixture):
@ -24,6 +26,52 @@ async def test_setup(hass: HomeAssistant, mock_entry: MockEntityFixture):
assert mock_entry.entry.unique_id == mock_entry.api.bootstrap.nvr.mac
async def test_setup_multiple(
hass: HomeAssistant,
mock_entry: MockEntityFixture,
mock_client,
mock_bootstrap: MockBootstrap,
):
"""Test working setup of unifiprotect entry."""
await hass.config_entries.async_setup(mock_entry.entry.entry_id)
await hass.async_block_till_done()
assert mock_entry.entry.state == ConfigEntryState.LOADED
assert mock_entry.api.update.called
assert mock_entry.entry.unique_id == mock_entry.api.bootstrap.nvr.mac
nvr = mock_bootstrap.nvr
nvr._api = mock_client
nvr.mac = "A1E00C826983"
nvr.id
mock_client.get_nvr = AsyncMock(return_value=nvr)
with patch("homeassistant.components.unifiprotect.ProtectApiClient") as mock_api:
mock_config = MockConfigEntry(
domain=DOMAIN,
data={
"host": "1.1.1.1",
"username": "test-username",
"password": "test-password",
"id": "UnifiProtect",
"port": 443,
"verify_ssl": False,
},
version=2,
)
mock_config.add_to_hass(hass)
mock_api.return_value = mock_client
await hass.config_entries.async_setup(mock_config.entry_id)
await hass.async_block_till_done()
assert mock_config.state == ConfigEntryState.LOADED
assert mock_client.update.called
assert mock_config.unique_id == mock_client.bootstrap.nvr.mac
async def test_reload(hass: HomeAssistant, mock_entry: MockEntityFixture):
"""Test updating entry reload entry."""

View file

@ -0,0 +1,146 @@
"""Test the UniFi Protect global services."""
# pylint: disable=protected-access
from __future__ import annotations
from unittest.mock import AsyncMock, Mock
import pytest
from pyunifiprotect.data import Light
from pyunifiprotect.exceptions import BadRequest
from homeassistant.components.unifiprotect.const import (
ATTR_MESSAGE,
DOMAIN,
SERVICE_ADD_DOORBELL_TEXT,
SERVICE_REMOVE_DOORBELL_TEXT,
SERVICE_SET_DEFAULT_DOORBELL_TEXT,
)
from homeassistant.const import ATTR_DEVICE_ID
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers import device_registry as dr
from .conftest import MockEntityFixture
@pytest.fixture(name="device")
async def device_fixture(hass: HomeAssistant, mock_entry: MockEntityFixture):
"""Fixture with entry setup to call services with."""
await hass.config_entries.async_setup(mock_entry.entry.entry_id)
await hass.async_block_till_done()
device_registry = await dr.async_get_registry(hass)
return list(device_registry.devices.values())[0]
@pytest.fixture(name="subdevice")
async def subdevice_fixture(
hass: HomeAssistant, mock_entry: MockEntityFixture, mock_light: Light
):
"""Fixture with entry setup to call services with."""
mock_light._api = mock_entry.api
mock_entry.api.bootstrap.lights = {
mock_light.id: mock_light,
}
await hass.config_entries.async_setup(mock_entry.entry.entry_id)
await hass.async_block_till_done()
device_registry = await dr.async_get_registry(hass)
return [d for d in device_registry.devices.values() if d.name != "UnifiProtect"][0]
async def test_global_service_bad_device(
hass: HomeAssistant, device: dr.DeviceEntry, mock_entry: MockEntityFixture
):
"""Test global service, invalid device ID."""
nvr = mock_entry.api.bootstrap.nvr
nvr.__fields__["add_custom_doorbell_message"] = Mock()
nvr.add_custom_doorbell_message = AsyncMock()
with pytest.raises(HomeAssistantError):
await hass.services.async_call(
DOMAIN,
SERVICE_ADD_DOORBELL_TEXT,
{ATTR_DEVICE_ID: "bad_device_id", ATTR_MESSAGE: "Test Message"},
blocking=True,
)
assert not nvr.add_custom_doorbell_message.called
async def test_global_service_exception(
hass: HomeAssistant, device: dr.DeviceEntry, mock_entry: MockEntityFixture
):
"""Test global service, unexpected error."""
nvr = mock_entry.api.bootstrap.nvr
nvr.__fields__["add_custom_doorbell_message"] = Mock()
nvr.add_custom_doorbell_message = AsyncMock(side_effect=BadRequest)
with pytest.raises(HomeAssistantError):
await hass.services.async_call(
DOMAIN,
SERVICE_ADD_DOORBELL_TEXT,
{ATTR_DEVICE_ID: device.id, ATTR_MESSAGE: "Test Message"},
blocking=True,
)
assert nvr.add_custom_doorbell_message.called
async def test_add_doorbell_text(
hass: HomeAssistant, device: dr.DeviceEntry, mock_entry: MockEntityFixture
):
"""Test add_doorbell_text service."""
nvr = mock_entry.api.bootstrap.nvr
nvr.__fields__["add_custom_doorbell_message"] = Mock()
nvr.add_custom_doorbell_message = AsyncMock()
await hass.services.async_call(
DOMAIN,
SERVICE_ADD_DOORBELL_TEXT,
{ATTR_DEVICE_ID: device.id, ATTR_MESSAGE: "Test Message"},
blocking=True,
)
nvr.add_custom_doorbell_message.assert_called_once_with("Test Message")
async def test_remove_doorbell_text(
hass: HomeAssistant, subdevice: dr.DeviceEntry, mock_entry: MockEntityFixture
):
"""Test remove_doorbell_text service."""
nvr = mock_entry.api.bootstrap.nvr
nvr.__fields__["remove_custom_doorbell_message"] = Mock()
nvr.remove_custom_doorbell_message = AsyncMock()
await hass.services.async_call(
DOMAIN,
SERVICE_REMOVE_DOORBELL_TEXT,
{ATTR_DEVICE_ID: subdevice.id, ATTR_MESSAGE: "Test Message"},
blocking=True,
)
nvr.remove_custom_doorbell_message.assert_called_once_with("Test Message")
async def test_set_default_doorbell_text(
hass: HomeAssistant, device: dr.DeviceEntry, mock_entry: MockEntityFixture
):
"""Test set_default_doorbell_text service."""
nvr = mock_entry.api.bootstrap.nvr
nvr.__fields__["set_default_doorbell_message"] = Mock()
nvr.set_default_doorbell_message = AsyncMock()
await hass.services.async_call(
DOMAIN,
SERVICE_SET_DEFAULT_DOORBELL_TEXT,
{ATTR_DEVICE_ID: device.id, ATTR_MESSAGE: "Test Message"},
blocking=True,
)
nvr.set_default_doorbell_message.assert_called_once_with("Test Message")