Add discovery support to unifiprotect (#64340)

This commit is contained in:
J. Nick Koston 2022-01-18 08:40:29 -10:00 committed by GitHub
parent 81461832c3
commit 88261c6c14
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
14 changed files with 450 additions and 26 deletions

View File

@ -33,6 +33,7 @@ from .const import (
PLATFORMS,
)
from .data import ProtectData
from .discovery import async_start_discovery
from .services import async_cleanup_services, async_setup_services
_LOGGER = logging.getLogger(__name__)
@ -43,6 +44,7 @@ SCAN_INTERVAL = timedelta(seconds=DEFAULT_SCAN_INTERVAL)
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Set up the UniFi Protect config entries."""
await async_start_discovery(hass)
session = async_create_clientsession(hass, cookie_jar=CookieJar(unsafe=True))
protect = ProtectApiClient(
host=entry.data[CONF_HOST],

View File

@ -10,6 +10,7 @@ from pyunifiprotect.data.nvr import NVR
import voluptuous as vol
from homeassistant import config_entries
from homeassistant.components import dhcp, ssdp
from homeassistant.const import (
CONF_HOST,
CONF_ID,
@ -21,6 +22,7 @@ from homeassistant.const import (
from homeassistant.core import callback
from homeassistant.data_entry_flow import FlowResult
from homeassistant.helpers.aiohttp_client import async_create_clientsession
from homeassistant.helpers.typing import DiscoveryInfoType
from .const import (
CONF_ALL_UPDATES,
@ -32,6 +34,8 @@ from .const import (
MIN_REQUIRED_PROTECT_V,
OUTDATED_LOG_MESSAGE,
)
from .discovery import async_start_discovery
from .services import _async_unifi_mac_from_hass
_LOGGER = logging.getLogger(__name__)
@ -44,8 +48,77 @@ class ProtectFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
def __init__(self) -> None:
"""Init the config flow."""
super().__init__()
self.entry: config_entries.ConfigEntry | None = None
self._discovered_device: dict[str, str] = {}
async def async_step_dhcp(self, discovery_info: dhcp.DhcpServiceInfo) -> FlowResult:
"""Handle discovery via dhcp."""
_LOGGER.debug("Starting discovery via: %s", discovery_info)
return await self._async_discovery_handoff()
async def async_step_ssdp(self, discovery_info: ssdp.SsdpServiceInfo) -> FlowResult:
"""Handle a discovered UniFi device."""
_LOGGER.debug("Starting discovery via: %s", discovery_info)
return await self._async_discovery_handoff()
async def _async_discovery_handoff(self) -> FlowResult:
"""Ensure discovery is active."""
# Discovery requires an additional check so we use
# SSDP and DHCP to tell us to start it so it only
# runs on networks where unifi devices are present.
await async_start_discovery(self.hass)
return self.async_abort(reason="discovery_started")
async def async_step_discovery(
self, discovery_info: DiscoveryInfoType
) -> FlowResult:
"""Handle discovery."""
self._discovered_device = discovery_info
mac = _async_unifi_mac_from_hass(discovery_info["mac"])
await self.async_set_unique_id(mac)
self._abort_if_unique_id_configured(
updates={CONF_HOST: discovery_info["ip_address"]}
)
return await self.async_step_discovery_confirm()
async def async_step_discovery_confirm(
self, user_input: dict[str, Any] | None = None
) -> FlowResult:
"""Confirm discovery."""
errors: dict[str, str] = {}
discovery_info = self._discovered_device
if user_input is not None:
user_input[CONF_HOST] = discovery_info["ip_address"]
user_input[CONF_PORT] = DEFAULT_PORT
nvr_data, errors = await self._async_get_nvr_data(user_input)
if nvr_data and not errors:
return self._async_create_entry(nvr_data.name, user_input)
placeholders = {
"name": discovery_info["hostname"]
or discovery_info["platform"]
or f"NVR {discovery_info['mac']}",
"ip_address": discovery_info["ip_address"],
}
self.context["title_placeholders"] = placeholders
user_input = user_input or {}
return self.async_show_form(
step_id="discovery_confirm",
description_placeholders=placeholders,
data_schema=vol.Schema(
{
vol.Required(
CONF_VERIFY_SSL,
default=user_input.get(CONF_VERIFY_SSL, DEFAULT_VERIFY_SSL),
): bool,
vol.Required(
CONF_USERNAME, default=user_input.get(CONF_USERNAME)
): str,
vol.Required(CONF_PASSWORD): str,
}
),
errors=errors,
)
@staticmethod
@callback

View File

@ -0,0 +1,66 @@
"""The unifiprotect integration discovery."""
from __future__ import annotations
import asyncio
from datetime import timedelta
import logging
from typing import Any
from unifi_discovery import AIOUnifiScanner, UnifiDevice, UnifiService
from homeassistant import config_entries
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.event import async_track_time_interval
from .const import DOMAIN
_LOGGER = logging.getLogger(__name__)
DISCOVERY = "discovery"
DISCOVERY_INTERVAL = timedelta(minutes=60)
async def async_start_discovery(hass: HomeAssistant) -> None:
"""Start discovery."""
domain_data = hass.data.setdefault(DOMAIN, {})
if DISCOVERY in domain_data:
return
domain_data[DISCOVERY] = True
async def _async_discovery(*_: Any) -> None:
async_trigger_discovery(hass, await async_discover_devices(hass))
# Do not block startup since discovery takes 31s or more
asyncio.create_task(_async_discovery())
async_track_time_interval(hass, _async_discovery, DISCOVERY_INTERVAL)
async def async_discover_devices(hass: HomeAssistant) -> list[UnifiDevice]:
"""Discover devices."""
scanner = AIOUnifiScanner()
devices = await scanner.async_scan()
_LOGGER.debug("Found devices: %s", devices)
return devices
@callback
def async_trigger_discovery(
hass: HomeAssistant,
discovered_devices: list[UnifiDevice],
) -> None:
"""Trigger config flows for discovered devices."""
for device in discovered_devices:
if device.services[UnifiService.Protect] and device.hw_addr:
hass.async_create_task(
hass.config_entries.flow.async_init(
DOMAIN,
context={"source": config_entries.SOURCE_DISCOVERY},
data={
"ip_address": device.source_ip,
"mac": device.hw_addr,
"hostname": device.hostname, # can be None
"platform": device.platform, # can be None
},
)
)

View File

@ -4,7 +4,7 @@
"config_flow": true,
"documentation": "https://www.home-assistant.io/integrations/unifiprotect",
"requirements": [
"pyunifiprotect==3.1.1"
"pyunifiprotect==3.1.1", "unifi-discovery==1.0.0"
],
"dependencies": [
"http"
@ -15,5 +15,48 @@
"@bdraco"
],
"quality_scale": "platinum",
"iot_class": "local_push"
"iot_class": "local_push",
"dhcp": [
{
"macaddress": "B4FBE4*"
},
{
"macaddress": "802AA8*"
},
{
"macaddress": "F09FC2*"
},
{
"macaddress": "68D79A*"
},
{
"macaddress": "18E829*"
},
{
"macaddress": "245A4C*"
},
{
"macaddress": "784558*"
},
{
"macaddress": "E063DA*"
},
{
"macaddress": "265A4C*"
}
],
"ssdp": [
{
"manufacturer": "Ubiquiti Networks",
"modelDescription": "UniFi Dream Machine"
},
{
"manufacturer": "Ubiquiti Networks",
"modelDescription": "UniFi Dream Machine Pro"
},
{
"manufacturer": "Ubiquiti Networks",
"modelDescription": "UniFi Dream Machine SE"
}
]
}

View File

@ -1,5 +1,6 @@
{
"config": {
"flow_title": "{name} ({ip_address})",
"step": {
"user": {
"title": "UniFi Protect Setup",
@ -19,7 +20,15 @@
"username": "[%key:common::config_flow::data::username%]",
"password": "[%key:common::config_flow::data::password%]"
}
}
},
"discovery_confirm": {
"description": "Do you want to setup {name} ({ip_address})?",
"data": {
"verify_ssl": "[%key:common::config_flow::data::verify_ssl%]",
"username": "[%key:common::config_flow::data::username%]",
"password": "[%key:common::config_flow::data::password%]"
}
}
},
"error": {
"cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
@ -27,7 +36,8 @@
"protect_version": "Minimum required version is v1.20.0. Please upgrade UniFi Protect and then retry."
},
"abort": {
"already_configured": "[%key:common::config_flow::abort::already_configured_device%]"
"already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
"discovery_started": "Discovery started"
}
},
"options": {

View File

@ -1,15 +1,24 @@
{
"config": {
"abort": {
"already_configured": "Device is already configured"
"already_configured": "Device is already configured",
"discovery_started": "Discovery started"
},
"error": {
"cannot_connect": "Failed to connect",
"invalid_auth": "Invalid authentication",
"protect_version": "Minimum required version is v1.20.0. Please upgrade UniFi Protect and then retry.",
"unknown": "Unexpected error"
"protect_version": "Minimum required version is v1.20.0. Please upgrade UniFi Protect and then retry."
},
"flow_title": "{name} ({ip_address})",
"step": {
"discovery_confirm": {
"data": {
"password": "Password",
"username": "Username",
"verify_ssl": "Verify SSL certificate"
},
"description": "Do you want to setup {name} ({ip_address})?"
},
"reauth_confirm": {
"data": {
"host": "IP/Host of UniFi Protect Server",

View File

@ -546,6 +546,42 @@ DHCP = [
"domain": "twinkly",
"hostname": "twinkly_*"
},
{
"domain": "unifiprotect",
"macaddress": "B4FBE4*"
},
{
"domain": "unifiprotect",
"macaddress": "802AA8*"
},
{
"domain": "unifiprotect",
"macaddress": "F09FC2*"
},
{
"domain": "unifiprotect",
"macaddress": "68D79A*"
},
{
"domain": "unifiprotect",
"macaddress": "18E829*"
},
{
"domain": "unifiprotect",
"macaddress": "245A4C*"
},
{
"domain": "unifiprotect",
"macaddress": "784558*"
},
{
"domain": "unifiprotect",
"macaddress": "E063DA*"
},
{
"domain": "unifiprotect",
"macaddress": "265A4C*"
},
{
"domain": "verisure",
"macaddress": "0023C1*"

View File

@ -238,6 +238,20 @@ SSDP = {
"modelDescription": "UniFi Dream Machine Pro"
}
],
"unifiprotect": [
{
"manufacturer": "Ubiquiti Networks",
"modelDescription": "UniFi Dream Machine"
},
{
"manufacturer": "Ubiquiti Networks",
"modelDescription": "UniFi Dream Machine Pro"
},
{
"manufacturer": "Ubiquiti Networks",
"modelDescription": "UniFi Dream Machine SE"
}
],
"upnp": [
{
"st": "urn:schemas-upnp-org:device:InternetGatewayDevice:1"

View File

@ -2391,6 +2391,9 @@ twilio==6.32.0
# homeassistant.components.rainforest_eagle
uEagle==0.0.2
# homeassistant.components.unifiprotect
unifi-discovery==1.0.0
# homeassistant.components.unifiled
unifiled==0.11

View File

@ -1449,6 +1449,9 @@ twilio==6.32.0
# homeassistant.components.rainforest_eagle
uEagle==0.0.2
# homeassistant.components.unifiprotect
unifi-discovery==1.0.0
# homeassistant.components.upb
upb_lib==0.4.12

View File

@ -1 +1,36 @@
"""Tests for the UniFi Protect integration."""
from contextlib import contextmanager
from unittest.mock import AsyncMock, MagicMock, patch
from unifi_discovery import AIOUnifiScanner, UnifiDevice, UnifiService
DEVICE_HOSTNAME = "unvr"
DEVICE_IP_ADDRESS = "127.0.0.1"
DEVICE_MAC_ADDRESS = "aa:bb:cc:dd:ee:ff"
UNIFI_DISCOVERY = UnifiDevice(
source_ip=DEVICE_IP_ADDRESS,
hw_addr=DEVICE_MAC_ADDRESS,
platform=DEVICE_HOSTNAME,
hostname=DEVICE_HOSTNAME,
services={UnifiService.Protect: True},
)
def _patch_discovery(device=None, no_device=False):
mock_aio_discovery = MagicMock(auto_spec=AIOUnifiScanner)
scanner_return = [] if no_device else [device or UNIFI_DISCOVERY]
mock_aio_discovery.async_scan = AsyncMock(return_value=scanner_return)
mock_aio_discovery.found_devices = scanner_return
@contextmanager
def _patcher():
with patch(
"homeassistant.components.unifiprotect.discovery.AIOUnifiScanner",
return_value=mock_aio_discovery,
):
yield
return _patcher()

View File

@ -23,6 +23,8 @@ from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.entity import EntityDescription
import homeassistant.util.dt as dt_util
from . import _patch_discovery
from tests.common import MockConfigEntry, async_fire_time_changed, load_fixture
MAC_ADDR = "aa:bb:cc:dd:ee:ff"
@ -73,6 +75,24 @@ def mock_nvr_fixture():
NVR.__config__.validate_assignment = True
@pytest.fixture(name="mock_ufp_config_entry")
def mock_ufp_config_entry():
"""Mock the unifiprotect config entry."""
return MockConfigEntry(
domain=DOMAIN,
data={
"host": "1.1.1.1",
"username": "test-username",
"password": "test-password",
"id": "UnifiProtect",
"port": 443,
"verify_ssl": False,
},
version=2,
)
@pytest.fixture(name="mock_old_nvr")
def mock_old_nvr_fixture():
"""Mock UniFi Protect Camera device."""
@ -122,28 +142,20 @@ def mock_client(mock_bootstrap: MockBootstrap):
@pytest.fixture
def mock_entry(
hass: HomeAssistant, mock_client # pylint: disable=redefined-outer-name
hass: HomeAssistant,
mock_ufp_config_entry: MockConfigEntry,
mock_client, # pylint: disable=redefined-outer-name
):
"""Mock ProtectApiClient for testing."""
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)
with _patch_discovery(no_device=True), patch(
"homeassistant.components.unifiprotect.ProtectApiClient"
) as mock_api:
mock_ufp_config_entry.add_to_hass(hass)
mock_api.return_value = mock_client
yield MockEntityFixture(mock_config, mock_client)
yield MockEntityFixture(mock_ufp_config_entry, mock_client)
@pytest.fixture

View File

@ -3,10 +3,12 @@ from __future__ import annotations
from unittest.mock import patch
import pytest
from pyunifiprotect import NotAuthorized, NvrError
from pyunifiprotect.data.nvr import NVR
from homeassistant import config_entries
from homeassistant.components import dhcp, ssdp
from homeassistant.components.unifiprotect.const import (
CONF_ALL_UPDATES,
CONF_DISABLE_RTSP,
@ -21,10 +23,35 @@ from homeassistant.data_entry_flow import (
)
from homeassistant.helpers import device_registry as dr
from . import DEVICE_HOSTNAME, DEVICE_IP_ADDRESS, DEVICE_MAC_ADDRESS, _patch_discovery
from .conftest import MAC_ADDR
from tests.common import MockConfigEntry
DHCP_DISCOVERY = dhcp.DhcpServiceInfo(
hostname=DEVICE_HOSTNAME,
ip=DEVICE_IP_ADDRESS,
macaddress=DEVICE_MAC_ADDRESS,
)
SSDP_DISCOVERY = (
ssdp.SsdpServiceInfo(
ssdp_usn="mock_usn",
ssdp_st="mock_st",
ssdp_location=f"http://{DEVICE_IP_ADDRESS}:41417/rootDesc.xml",
upnp={
"friendlyName": "UniFi Dream Machine",
"modelDescription": "UniFi Dream Machine Pro",
"serialNumber": DEVICE_MAC_ADDRESS,
},
),
)
UNIFI_DISCOVERY_DICT = {
"ip_address": DEVICE_IP_ADDRESS,
"mac": DEVICE_MAC_ADDRESS,
"hostname": DEVICE_HOSTNAME,
"platform": DEVICE_HOSTNAME,
}
async def test_form(hass: HomeAssistant, mock_nvr: NVR) -> None:
"""Test we get the form."""
@ -208,7 +235,9 @@ async def test_form_options(hass: HomeAssistant, mock_client) -> None:
)
mock_config.add_to_hass(hass)
with patch("homeassistant.components.unifiprotect.ProtectApiClient") as mock_api:
with _patch_discovery(), patch(
"homeassistant.components.unifiprotect.ProtectApiClient"
) as mock_api:
mock_api.return_value = mock_client
await hass.config_entries.async_setup(mock_config.entry_id)
@ -231,3 +260,73 @@ async def test_form_options(hass: HomeAssistant, mock_client) -> None:
"disable_rtsp": True,
"override_connection_host": True,
}
@pytest.mark.parametrize(
"source, data",
[
(config_entries.SOURCE_DHCP, DHCP_DISCOVERY),
(config_entries.SOURCE_SSDP, SSDP_DISCOVERY),
],
)
async def test_discovered_by_ssdp_or_dhcp(
hass: HomeAssistant, source: str, data: dhcp.DhcpServiceInfo | ssdp.SsdpServiceInfo
) -> None:
"""Test we handoff to unifi-discovery when discovered via ssdp or dhcp."""
with _patch_discovery():
result = await hass.config_entries.flow.async_init(
DOMAIN,
context={"source": source},
data=data,
)
await hass.async_block_till_done()
assert result["type"] == RESULT_TYPE_ABORT
assert result["reason"] == "discovery_started"
async def test_discovered_by_unifi_discovery(
hass: HomeAssistant, mock_nvr: NVR
) -> None:
"""Test a discovery from unifi-discovery."""
with _patch_discovery():
result = await hass.config_entries.flow.async_init(
DOMAIN,
context={"source": config_entries.SOURCE_DISCOVERY},
data=UNIFI_DISCOVERY_DICT,
)
await hass.async_block_till_done()
assert result["type"] == RESULT_TYPE_FORM
assert result["step_id"] == "discovery_confirm"
assert not result["errors"]
with patch(
"homeassistant.components.unifiprotect.config_flow.ProtectApiClient.get_nvr",
return_value=mock_nvr,
), patch(
"homeassistant.components.unifiprotect.async_setup_entry",
return_value=True,
) as mock_setup_entry:
result2 = await hass.config_entries.flow.async_configure(
result["flow_id"],
{
"username": "test-username",
"password": "test-password",
},
)
await hass.async_block_till_done()
assert result2["type"] == RESULT_TYPE_CREATE_ENTRY
assert result2["title"] == "UnifiProtect"
assert result2["data"] == {
"host": DEVICE_IP_ADDRESS,
"username": "test-username",
"password": "test-password",
"id": "UnifiProtect",
"port": 443,
"verify_ssl": False,
}
assert len(mock_setup_entry.mock_calls) == 1

View File

@ -7,9 +7,10 @@ from pyunifiprotect import NotAuthorized, NvrError
from pyunifiprotect.data import NVR
from homeassistant.components.unifiprotect.const import CONF_DISABLE_RTSP, DOMAIN
from homeassistant.config_entries import ConfigEntryState
from homeassistant.config_entries import ConfigEntry, ConfigEntryState
from homeassistant.core import HomeAssistant
from . import _patch_discovery
from .conftest import MockBootstrap, MockEntityFixture
from tests.common import MockConfigEntry
@ -156,3 +157,21 @@ async def test_setup_failed_auth(hass: HomeAssistant, mock_entry: MockEntityFixt
await hass.config_entries.async_setup(mock_entry.entry.entry_id)
assert mock_entry.entry.state == ConfigEntryState.SETUP_ERROR
assert not mock_entry.api.update.called
async def test_setup_starts_discovery(
hass: HomeAssistant, mock_ufp_config_entry: ConfigEntry, mock_client
):
"""Test setting up will start discovery."""
with _patch_discovery(), patch(
"homeassistant.components.unifiprotect.ProtectApiClient"
) as mock_api:
mock_ufp_config_entry.add_to_hass(hass)
mock_api.return_value = mock_client
mock_entry = MockEntityFixture(mock_ufp_config_entry, mock_client)
await hass.config_entries.async_setup(mock_entry.entry.entry_id)
await hass.async_block_till_done()
assert mock_entry.entry.state == ConfigEntryState.LOADED
await hass.async_block_till_done()
assert len(hass.config_entries.flow.async_progress_by_handler(DOMAIN)) == 1