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

Make UniFi utilise forward_entry_setups (#74835)

Browse Source
This commit is contained in:
Robert Svensson 2022-07-18 17:39:38 +02:00 committed by GitHub
parent b3ef6f4d04
commit 3144d179e0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 103 additions and 126 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

50
homeassistant/components/unifi/__init__.py
View File

@ -5,19 +5,13 @@ from typing import Any
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import EVENT_HOMEASSISTANT_STOP
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import device_registry as dr
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
from homeassistant.helpers.storage import Store
from homeassistant.helpers.typing import ConfigType
from .const import (
ATTR_MANUFACTURER,
CONF_CONTROLLER,
DOMAIN as UNIFI_DOMAIN,
LOGGER,
UNIFI_WIRELESS_CLIENTS,
)
from .controller import UniFiController
from .const import CONF_CONTROLLER, DOMAIN as UNIFI_DOMAIN, UNIFI_WIRELESS_CLIENTS
from .controller import PLATFORMS, UniFiController, get_unifi_controller
from .errors import AuthenticationRequired, CannotConnect
from .services import async_setup_services, async_unload_services
SAVE_DELAY = 10
@ -40,9 +34,16 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
# Flat configuration was introduced with 2021.3
await async_flatten_entry_data(hass, config_entry)
controller = UniFiController(hass, config_entry)
if not await controller.async_setup():
return False
try:
api = await get_unifi_controller(hass, config_entry.data)
controller = UniFiController(hass, config_entry, api)
await controller.initialize()
except CannotConnect as err:
raise ConfigEntryNotReady from err
except AuthenticationRequired as err:
raise ConfigEntryAuthFailed from err
# Unique ID was introduced with 2021.3
if config_entry.unique_id is None:
@ -50,30 +51,19 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
config_entry, unique_id=controller.site_id
)
if not hass.data[UNIFI_DOMAIN]:
hass.data[UNIFI_DOMAIN][config_entry.entry_id] = controller
await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
await controller.async_update_device_registry()
if len(hass.data[UNIFI_DOMAIN]) == 1:
async_setup_services(hass)
hass.data[UNIFI_DOMAIN][config_entry.entry_id] = controller
api.start_websocket()
config_entry.async_on_unload(
hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, controller.shutdown)
)
LOGGER.debug("UniFi Network config options %s", config_entry.options)
if controller.mac is None:
return True
device_registry = dr.async_get(hass)
device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id,
configuration_url=controller.api.url,
connections={(CONNECTION_NETWORK_MAC, controller.mac)},
default_manufacturer=ATTR_MANUFACTURER,
default_model="UniFi Network",
default_name="UniFi Network",
)
return True

14
homeassistant/components/unifi/config_flow.py
View File

@ -9,6 +9,7 @@ from __future__ import annotations
from collections.abc import Mapping
import socket
from types import MappingProxyType
from typing import Any
from urllib.parse import urlparse
@ -46,7 +47,7 @@ from .const import (
DEFAULT_POE_CLIENTS,
DOMAIN as UNIFI_DOMAIN,
)
from .controller import UniFiController, get_controller
from .controller import UniFiController, get_unifi_controller
from .errors import AuthenticationRequired, CannotConnect
DEFAULT_PORT = 443
@ -99,16 +100,9 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
}
try:
controller = await get_controller(
self.hass,
host=self.config[CONF_HOST],
username=self.config[CONF_USERNAME],
password=self.config[CONF_PASSWORD],
port=self.config[CONF_PORT],
site=self.config[CONF_SITE_ID],
verify_ssl=self.config[CONF_VERIFY_SSL],
controller = await get_unifi_controller(
self.hass, MappingProxyType(self.config)
)
sites = await controller.sites()
except AuthenticationRequired:

94
homeassistant/components/unifi/controller.py
View File

@ -4,6 +4,8 @@ from __future__ import annotations
import asyncio
from datetime import datetime, timedelta
import ssl
from types import MappingProxyType
from typing import Any
from aiohttp import CookieJar
import aiounifi
@ -36,14 +38,19 @@ from homeassistant.const import (
Platform,
)
from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
from homeassistant.helpers import aiohttp_client, entity_registry as er
from homeassistant.helpers import (
aiohttp_client,
device_registry as dr,
entity_registry as er,
)
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
from homeassistant.helpers.dispatcher import async_dispatcher_send
from homeassistant.helpers.entity_registry import async_entries_for_config_entry
from homeassistant.helpers.event import async_track_time_interval
import homeassistant.util.dt as dt_util
from .const import (
ATTR_MANUFACTURER,
CONF_ALLOW_BANDWIDTH_SENSORS,
CONF_ALLOW_UPTIME_SENSORS,
CONF_BLOCK_CLIENT,
@ -91,12 +98,15 @@ DEVICE_CONNECTED = (
class UniFiController:
"""Manages a single UniFi Network instance."""
def __init__(self, hass, config_entry):
def __init__(self, hass, config_entry, api):
"""Initialize the system."""
self.hass = hass
self.config_entry = config_entry
self.api = api
api.callback = self.async_unifi_signalling_callback
self.available = True
self.api = None
self.progress = None
self.wireless_clients = None
@ -295,36 +305,18 @@ class UniFiController:
unifi_wireless_clients = self.hass.data[UNIFI_WIRELESS_CLIENTS]
unifi_wireless_clients.update_data(self.wireless_clients, self.config_entry)
async def async_setup(self):
async def initialize(self):
"""Set up a UniFi Network instance."""
try:
self.api = await get_controller(
self.hass,
host=self.config_entry.data[CONF_HOST],
username=self.config_entry.data[CONF_USERNAME],
password=self.config_entry.data[CONF_PASSWORD],
port=self.config_entry.data[CONF_PORT],
site=self.config_entry.data[CONF_SITE_ID],
verify_ssl=self.config_entry.data[CONF_VERIFY_SSL],
async_callback=self.async_unifi_signalling_callback,
)
await self.api.initialize()
sites = await self.api.sites()
description = await self.api.site_description()
except CannotConnect as err:
raise ConfigEntryNotReady from err
except AuthenticationRequired as err:
raise ConfigEntryAuthFailed from err
await self.api.initialize()
sites = await self.api.sites()
for site in sites.values():
if self.site == site["name"]:
self.site_id = site["_id"]
self._site_name = site["desc"]
break
description = await self.api.site_description()
self._site_role = description[0]["site_role"]
# Restore clients that are not a part of active clients list.
@ -357,18 +349,12 @@ class UniFiController:
self.wireless_clients = wireless_clients.get_data(self.config_entry)
self.update_wireless_clients()
self.hass.config_entries.async_setup_platforms(self.config_entry, PLATFORMS)
self.api.start_websocket()
self.config_entry.add_update_listener(self.async_config_entry_updated)
self._cancel_heartbeat_check = async_track_time_interval(
self.hass, self._async_check_for_stale, CHECK_HEARTBEAT_INTERVAL
)
return True
@callback
def async_heartbeat(
self, unique_id: str, heartbeat_expire_time: datetime | None = None
@ -397,6 +383,22 @@ class UniFiController:
for unique_id in unique_ids_to_remove:
del self._heartbeat_time[unique_id]
async def async_update_device_registry(self) -> None:
"""Update device registry."""
if self.mac is None:
return
device_registry = dr.async_get(self.hass)
device_registry.async_get_or_create(
config_entry_id=self.config_entry.entry_id,
configuration_url=self.api.url,
connections={(CONNECTION_NETWORK_MAC, self.mac)},
default_manufacturer=ATTR_MANUFACTURER,
default_model="UniFi Network",
default_name="UniFi Network",
)
@staticmethod
async def async_config_entry_updated(
hass: HomeAssistant, config_entry: ConfigEntry
@ -463,13 +465,14 @@ class UniFiController:
return True
async def get_controller(
hass, host, username, password, port, site, verify_ssl, async_callback=None
):
async def get_unifi_controller(
hass: HomeAssistant,
config: MappingProxyType[str, Any],
) -> aiounifi.Controller:
"""Create a controller object and verify authentication."""
sslcontext = None
if verify_ssl:
if verify_ssl := bool(config.get(CONF_VERIFY_SSL)):
session = aiohttp_client.async_get_clientsession(hass)
if isinstance(verify_ssl, str):
sslcontext = ssl.create_default_context(cafile=verify_ssl)
@ -479,14 +482,13 @@ async def get_controller(
)
controller = aiounifi.Controller(
host,
username=username,
password=password,
port=port,
site=site,
host=config[CONF_HOST],
username=config[CONF_USERNAME],
password=config[CONF_PASSWORD],
port=config[CONF_PORT],
site=config[CONF_SITE_ID],
websession=session,
sslcontext=sslcontext,
callback=async_callback,
)
try:
@ -498,7 +500,7 @@ async def get_controller(
except aiounifi.Unauthorized as err:
LOGGER.warning(
"Connected to UniFi Network at %s but not registered: %s",
host,
config[CONF_HOST],
err,
)
raise AuthenticationRequired from err
@ -510,13 +512,15 @@ async def get_controller(
aiounifi.RequestError,
aiounifi.ResponseError,
) as err:
LOGGER.error("Error connecting to the UniFi Network at %s: %s", host, err)
LOGGER.error(
"Error connecting to the UniFi Network at %s: %s", config[CONF_HOST], err
)
raise CannotConnect from err
except aiounifi.LoginRequired as err:
LOGGER.warning(
"Connected to UniFi Network at %s but login required: %s",
host,
config[CONF_HOST],
err,
)
raise AuthenticationRequired from err

24
tests/components/unifi/test_controller.py
View File

@ -30,7 +30,7 @@ from homeassistant.components.unifi.const import (
from homeassistant.components.unifi.controller import (
PLATFORMS,
RETRY_TIMER,
get_controller,
get_unifi_controller,
)
from homeassistant.components.unifi.errors import AuthenticationRequired, CannotConnect
from homeassistant.const import (
@ -271,7 +271,7 @@ async def test_controller_mac(hass, aioclient_mock):
async def test_controller_not_accessible(hass):
"""Retry to login gets scheduled when connection fails."""
with patch(
"homeassistant.components.unifi.controller.get_controller",
"homeassistant.components.unifi.controller.get_unifi_controller",
side_effect=CannotConnect,
):
await setup_unifi_integration(hass)
@ -281,7 +281,7 @@ async def test_controller_not_accessible(hass):
async def test_controller_trigger_reauth_flow(hass):
"""Failed authentication trigger a reauthentication flow."""
with patch(
"homeassistant.components.unifi.controller.get_controller",
"homeassistant.components.unifi.get_unifi_controller",
side_effect=AuthenticationRequired,
), patch.object(hass.config_entries.flow, "async_init") as mock_flow_init:
await setup_unifi_integration(hass)
@ -292,7 +292,7 @@ async def test_controller_trigger_reauth_flow(hass):
async def test_controller_unknown_error(hass):
"""Unknown errors are handled."""
with patch(
"homeassistant.components.unifi.controller.get_controller",
"homeassistant.components.unifi.controller.get_unifi_controller",
side_effect=Exception,
):
await setup_unifi_integration(hass)
@ -470,22 +470,22 @@ async def test_reconnect_mechanism_exceptions(
mock_reconnect.assert_called_once()
async def test_get_controller(hass):
async def test_get_unifi_controller(hass):
"""Successful call."""
with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(
"aiounifi.Controller.login", return_value=True
):
assert await get_controller(hass, **CONTROLLER_DATA)
assert await get_unifi_controller(hass, CONTROLLER_DATA)
async def test_get_controller_verify_ssl_false(hass):
async def test_get_unifi_controller_verify_ssl_false(hass):
"""Successful call with verify ssl set to false."""
controller_data = dict(CONTROLLER_DATA)
controller_data[CONF_VERIFY_SSL] = False
with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(
"aiounifi.Controller.login", return_value=True
):
assert await get_controller(hass, **controller_data)
assert await get_unifi_controller(hass, controller_data)
@pytest.mark.parametrize(
@ -501,9 +501,11 @@ async def test_get_controller_verify_ssl_false(hass):
(aiounifi.AiounifiException, AuthenticationRequired),
],
)
async def test_get_controller_fails_to_connect(hass, side_effect, raised_exception):
"""Check that get_controller can handle controller being unavailable."""
async def test_get_unifi_controller_fails_to_connect(
hass, side_effect, raised_exception
):
"""Check that get_unifi_controller can handle controller being unavailable."""
with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(
"aiounifi.Controller.login", side_effect=side_effect
), pytest.raises(raised_exception):
await get_controller(hass, **CONTROLLER_DATA)
await get_unifi_controller(hass, CONTROLLER_DATA)

47
tests/components/unifi/test_init.py
View File

@ -1,10 +1,10 @@
"""Test UniFi Network integration setup process."""
from unittest.mock import AsyncMock, patch
from unittest.mock import patch
from homeassistant.components import unifi
from homeassistant.components.unifi import async_flatten_entry_data
from homeassistant.components.unifi.const import CONF_CONTROLLER, DOMAIN as UNIFI_DOMAIN
from homeassistant.helpers import device_registry as dr
from homeassistant.components.unifi.errors import AuthenticationRequired, CannotConnect
from homeassistant.setup import async_setup_component
from .test_controller import (
@ -29,40 +29,27 @@ async def test_successful_config_entry(hass, aioclient_mock):
assert hass.data[UNIFI_DOMAIN]
async def test_controller_fail_setup(hass):
"""Test that a failed setup still stores controller."""
with patch("homeassistant.components.unifi.UniFiController") as mock_controller:
mock_controller.return_value.async_setup = AsyncMock(return_value=False)
async def test_setup_entry_fails_config_entry_not_ready(hass):
"""Failed authentication trigger a reauthentication flow."""
with patch(
"homeassistant.components.unifi.get_unifi_controller",
side_effect=CannotConnect,
):
await setup_unifi_integration(hass)
assert hass.data[UNIFI_DOMAIN] == {}
async def test_controller_mac(hass):
"""Test that configured options for a host are loaded via config entry."""
entry = MockConfigEntry(
domain=UNIFI_DOMAIN, data=ENTRY_CONFIG, unique_id="1", entry_id=1
)
entry.add_to_hass(hass)
async def test_setup_entry_fails_trigger_reauth_flow(hass):
"""Failed authentication trigger a reauthentication flow."""
with patch(
"homeassistant.components.unifi.get_unifi_controller",
side_effect=AuthenticationRequired,
), patch.object(hass.config_entries.flow, "async_init") as mock_flow_init:
await setup_unifi_integration(hass)
mock_flow_init.assert_called_once()
with patch("homeassistant.components.unifi.UniFiController") as mock_controller:
mock_controller.return_value.async_setup = AsyncMock(return_value=True)
mock_controller.return_value.mac = "mac1"
mock_controller.return_value.api.url = "https://123:443"
assert await unifi.async_setup_entry(hass, entry) is True
assert len(mock_controller.mock_calls) == 2
device_registry = dr.async_get(hass)
device = device_registry.async_get_or_create(
config_entry_id=entry.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, "mac1")},
)
assert device.configuration_url == "https://123:443"
assert device.manufacturer == "Ubiquiti Networks"
assert device.model == "UniFi Network"
assert device.name == "UniFi Network"
assert device.sw_version is None
assert hass.data[UNIFI_DOMAIN] == {}
async def test_flatten_entry_data(hass):