From 3144d179e0073f780c8d6738457f9e455c641339 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Mon, 18 Jul 2022 17:39:38 +0200 Subject: [PATCH] Make UniFi utilise forward_entry_setups (#74835) --- homeassistant/components/unifi/__init__.py | 50 ++++------ homeassistant/components/unifi/config_flow.py | 14 +-- homeassistant/components/unifi/controller.py | 94 ++++++++++--------- tests/components/unifi/test_controller.py | 24 ++--- tests/components/unifi/test_init.py | 47 ++++------ 5 files changed, 103 insertions(+), 126 deletions(-) diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index 1369bb69e1b..086bae8d8cf 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/homeassistant/components/unifi/__init__.py @@ -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 diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index 2f49c15e4d8..4944dd91296 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/homeassistant/components/unifi/config_flow.py @@ -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: diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index fa92568b477..7446d6abbff 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -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 diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index 625afbb4ec6..e420d031f46 100644 --- a/tests/components/unifi/test_controller.py +++ b/tests/components/unifi/test_controller.py @@ -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) diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py index f183e1c22ff..03ea89097c5 100644 --- a/tests/components/unifi/test_init.py +++ b/tests/components/unifi/test_init.py @@ -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):