From 83a75b02ea84fe34726c5215044ac9a185c36487 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Mon, 1 Feb 2021 17:55:16 +0100 Subject: [PATCH] Code quality improvements to UniFi integration (#45794) * Do less inside try statements * Replace controller id with config entry id since it doesn't serve a purpose anymore * Improve how controller connection state is communicated in the client and device tracker Remove the need to disable arguments-differ lint * Remove broad exception handling from config flow I'm not sure there ever was a reason for this more than to catch all exceptions * Replace site string with constant for SSDP title_placeholders * Unload platforms in the defacto way * Noone reads the method descriptions * Improve file descriptions --- homeassistant/components/unifi/__init__.py | 14 +--- homeassistant/components/unifi/config_flow.py | 67 ++++++++----------- homeassistant/components/unifi/const.py | 2 - homeassistant/components/unifi/controller.py | 53 ++++++++------- .../components/unifi/device_tracker.py | 24 ++++--- homeassistant/components/unifi/sensor.py | 6 +- homeassistant/components/unifi/switch.py | 7 +- tests/components/unifi/test_config_flow.py | 26 ------- tests/components/unifi/test_controller.py | 8 ++- 9 files changed, 89 insertions(+), 118 deletions(-) diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index 439073497a2b..a9d39251838c 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/homeassistant/components/unifi/__init__.py @@ -1,9 +1,8 @@ -"""Support for devices connected to UniFi POE.""" +"""Integration to UniFi controllers and its various features.""" from homeassistant.const import EVENT_HOMEASSISTANT_STOP from homeassistant.core import callback from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC -from .config_flow import get_controller_id_from_config_entry from .const import ( ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN, @@ -82,21 +81,12 @@ class UnifiWirelessClients: @callback def get_data(self, config_entry): """Get data related to a specific controller.""" - controller_id = get_controller_id_from_config_entry(config_entry) - key = config_entry.entry_id - if controller_id in self.data: - key = controller_id - - data = self.data.get(key, {"wireless_devices": []}) + data = self.data.get(config_entry.entry_id, {"wireless_devices": []}) return set(data["wireless_devices"]) @callback def update_data(self, data, config_entry): """Update data and schedule to save to file.""" - controller_id = get_controller_id_from_config_entry(config_entry) - if controller_id in self.data: - self.data.pop(controller_id) - self.data[config_entry.entry_id] = {"wireless_devices": list(data)} self._store.async_delay_save(self._data_to_save, SAVE_DELAY) diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index 85fe55a4076b..07b81621750f 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/homeassistant/components/unifi/config_flow.py @@ -1,4 +1,10 @@ -"""Config flow for UniFi.""" +"""Config flow for UniFi. + +Provides user initiated configuration flow. +Discovery of controllers hosted on UDM and UDM Pro devices through SSDP. +Reauthentication when issue with credentials are reported. +Configuration of options through options flow. +""" import socket from urllib.parse import urlparse @@ -31,11 +37,9 @@ from .const import ( CONF_TRACK_CLIENTS, CONF_TRACK_DEVICES, CONF_TRACK_WIRED_CLIENTS, - CONTROLLER_ID, DEFAULT_DPI_RESTRICTIONS, DEFAULT_POE_CLIENTS, DOMAIN as UNIFI_DOMAIN, - LOGGER, ) from .controller import get_controller from .errors import AuthenticationRequired, CannotConnect @@ -51,15 +55,6 @@ MODEL_PORTS = { } -@callback -def get_controller_id_from_config_entry(config_entry): - """Return controller with a matching bridge id.""" - return CONTROLLER_ID.format( - host=config_entry.data[CONF_CONTROLLER][CONF_HOST], - site=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID], - ) - - class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN): """Handle a UniFi config flow.""" @@ -86,19 +81,26 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN): if user_input is not None: + self.config = { + CONF_HOST: user_input[CONF_HOST], + CONF_USERNAME: user_input[CONF_USERNAME], + CONF_PASSWORD: user_input[CONF_PASSWORD], + CONF_PORT: user_input.get(CONF_PORT), + CONF_VERIFY_SSL: user_input.get(CONF_VERIFY_SSL), + CONF_SITE_ID: DEFAULT_SITE_ID, + } + try: - self.config = { - CONF_HOST: user_input[CONF_HOST], - CONF_USERNAME: user_input[CONF_USERNAME], - CONF_PASSWORD: user_input[CONF_PASSWORD], - CONF_PORT: user_input.get(CONF_PORT), - CONF_VERIFY_SSL: user_input.get(CONF_VERIFY_SSL), - CONF_SITE_ID: DEFAULT_SITE_ID, - } - controller = await get_controller(self.hass, **self.config) - sites = await controller.sites() + + except AuthenticationRequired: + errors["base"] = "faulty_credentials" + + except CannotConnect: + errors["base"] = "service_unavailable" + + else: self.sites = {site["name"]: site["desc"] for site in sites.values()} if self.reauth_config.get(CONF_SITE_ID) in self.sites: @@ -108,19 +110,6 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN): return await self.async_step_site() - except AuthenticationRequired: - errors["base"] = "faulty_credentials" - - except CannotConnect: - errors["base"] = "service_unavailable" - - except Exception: # pylint: disable=broad-except - LOGGER.error( - "Unknown error connecting with UniFi Controller at %s", - user_input[CONF_HOST], - ) - return self.async_abort(reason="unknown") - host = self.config.get(CONF_HOST) if not host and await async_discover_unifi(self.hass): host = "unifi" @@ -214,7 +203,7 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN): return await self.async_step_user() async def async_step_ssdp(self, discovery_info): - """Handle a discovered unifi device.""" + """Handle a discovered UniFi device.""" parsed_url = urlparse(discovery_info[ssdp.ATTR_SSDP_LOCATION]) model_description = discovery_info[ssdp.ATTR_UPNP_MODEL_DESCRIPTION] mac_address = format_mac(discovery_info[ssdp.ATTR_UPNP_SERIAL]) @@ -232,7 +221,7 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN): # pylint: disable=no-member self.context["title_placeholders"] = { CONF_HOST: self.config[CONF_HOST], - CONF_SITE_ID: "default", + CONF_SITE_ID: DEFAULT_SITE_ID, } port = MODEL_PORTS.get(model_description) @@ -242,7 +231,7 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN): return await self.async_step_user() def _host_already_configured(self, host): - """See if we already have a unifi entry matching the host.""" + """See if we already have a UniFi entry matching the host.""" for entry in self._async_current_entries(): if not entry.data or CONF_CONTROLLER not in entry.data: continue @@ -271,7 +260,7 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow): return await self.async_step_simple_options() async def async_step_simple_options(self, user_input=None): - """For simple Jack.""" + """For users without advanced settings enabled.""" if user_input is not None: self.options.update(user_input) return await self._update_options() diff --git a/homeassistant/components/unifi/const.py b/homeassistant/components/unifi/const.py index ba16612a9034..94e2fad35edd 100644 --- a/homeassistant/components/unifi/const.py +++ b/homeassistant/components/unifi/const.py @@ -4,8 +4,6 @@ import logging LOGGER = logging.getLogger(__package__) DOMAIN = "unifi" -CONTROLLER_ID = "{host}-{site}" - CONF_CONTROLLER = "controller" CONF_SITE_ID = "site" diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 11e02d60a3f2..bd55f4479fad 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -51,7 +51,6 @@ from .const import ( CONF_TRACK_CLIENTS, CONF_TRACK_DEVICES, CONF_TRACK_WIRED_CLIENTS, - CONTROLLER_ID, DEFAULT_ALLOW_BANDWIDTH_SENSORS, DEFAULT_ALLOW_UPTIME_SENSORS, DEFAULT_DETECTION_TIME, @@ -109,9 +108,10 @@ class UniFiController: def load_config_entry_options(self): """Store attributes to avoid property call overhead since they are called frequently.""" - # Device tracker options options = self.config_entry.options + # Device tracker options + # Config entry option to not track clients. self.option_track_clients = options.get( CONF_TRACK_CLIENTS, DEFAULT_TRACK_CLIENTS @@ -157,11 +157,6 @@ class UniFiController: CONF_ALLOW_UPTIME_SENSORS, DEFAULT_ALLOW_UPTIME_SENSORS ) - @property - def controller_id(self): - """Return the controller ID.""" - return CONTROLLER_ID.format(host=self.host, site=self.site) - @property def host(self): """Return the host of this controller.""" @@ -260,25 +255,25 @@ class UniFiController: @property def signal_reachable(self) -> str: """Integration specific event to signal a change in connection status.""" - return f"unifi-reachable-{self.controller_id}" + return f"unifi-reachable-{self.config_entry.entry_id}" @property - def signal_update(self): + def signal_update(self) -> str: """Event specific per UniFi entry to signal new data.""" - return f"unifi-update-{self.controller_id}" + return f"unifi-update-{self.config_entry.entry_id}" @property - def signal_remove(self): + def signal_remove(self) -> str: """Event specific per UniFi entry to signal removal of entities.""" - return f"unifi-remove-{self.controller_id}" + return f"unifi-remove-{self.config_entry.entry_id}" @property - def signal_options_update(self): + def signal_options_update(self) -> str: """Event specific per UniFi entry to signal new options.""" - return f"unifi-options-{self.controller_id}" + return f"unifi-options-{self.config_entry.entry_id}" @property - def signal_heartbeat_missed(self): + def signal_heartbeat_missed(self) -> str: """Event specific per UniFi device tracker to signal new heartbeat missed.""" return "unifi-heartbeat-missed" @@ -309,14 +304,7 @@ class UniFiController: await self.api.initialize() sites = await self.api.sites() - - for site in sites.values(): - if self.site == site["name"]: - self._site_name = site["desc"] - break - description = await self.api.site_description() - self._site_role = description[0]["site_role"] except CannotConnect as err: raise ConfigEntryNotReady from err @@ -331,6 +319,13 @@ class UniFiController: ) return False + for site in sites.values(): + if self.site == site["name"]: + self._site_name = site["desc"] + break + + self._site_role = description[0]["site_role"] + # Restore clients that is not a part of active clients list. entity_registry = await self.hass.helpers.entity_registry.async_get_registry() for entity in entity_registry.entities.values(): @@ -452,10 +447,18 @@ class UniFiController: """ self.api.stop_websocket() - for platform in SUPPORTED_PLATFORMS: - await self.hass.config_entries.async_forward_entry_unload( - self.config_entry, platform + unload_ok = all( + await asyncio.gather( + *[ + self.hass.config_entries.async_forward_entry_unload( + self.config_entry, platform + ) + for platform in SUPPORTED_PLATFORMS + ] ) + ) + if not unload_ok: + return False for unsub_dispatcher in self.listeners: unsub_dispatcher() diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index 6a4d986d5b2a..ac28f7475f65 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -1,4 +1,4 @@ -"""Track devices using UniFi controllers.""" +"""Track both clients and devices using UniFi controllers.""" from datetime import timedelta from aiounifi.api import SOURCE_DATA, SOURCE_EVENT @@ -145,6 +145,7 @@ class UniFiClientTracker(UniFiClient, ScannerEntity): self.heartbeat_check = False self._is_connected = False + self._controller_connection_state_changed = False if client.last_seen: self._is_connected = ( @@ -175,14 +176,16 @@ class UniFiClientTracker(UniFiClient, ScannerEntity): @callback def async_signal_reachable_callback(self) -> None: """Call when controller connection state change.""" - self.async_update_callback(controller_state_change=True) + self._controller_connection_state_changed = True + super().async_signal_reachable_callback() - # pylint: disable=arguments-differ @callback - def async_update_callback(self, controller_state_change: bool = False) -> None: + def async_update_callback(self) -> None: """Update the clients state.""" - if controller_state_change: + if self._controller_connection_state_changed: + self._controller_connection_state_changed = False + if self.controller.available: self.schedule_update = True @@ -304,6 +307,7 @@ class UniFiDeviceTracker(UniFiBase, ScannerEntity): self.device = self._item self._is_connected = device.state == 1 + self._controller_connection_state_changed = False self.schedule_update = False async def async_added_to_hass(self) -> None: @@ -325,14 +329,16 @@ class UniFiDeviceTracker(UniFiBase, ScannerEntity): @callback def async_signal_reachable_callback(self) -> None: """Call when controller connection state change.""" - self.async_update_callback(controller_state_change=True) + self._controller_connection_state_changed = True + super().async_signal_reachable_callback() - # pylint: disable=arguments-differ @callback - def async_update_callback(self, controller_state_change: bool = False) -> None: + def async_update_callback(self) -> None: """Update the devices' state.""" - if controller_state_change: + if self._controller_connection_state_changed: + self._controller_connection_state_changed = False + if self.controller.available: if self._is_connected: self.schedule_update = True diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index c0b8cea09c21..f78ec614da15 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/homeassistant/components/unifi/sensor.py @@ -1,4 +1,8 @@ -"""Support for bandwidth sensors with UniFi clients.""" +"""Sensor platform for UniFi integration. + +Support for bandwidth sensors of network clients. +Support for uptime sensors of network clients. +""" from homeassistant.components.sensor import DEVICE_CLASS_TIMESTAMP, DOMAIN from homeassistant.const import DATA_MEGABYTES from homeassistant.core import callback diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 6aa42b0d2918..e596e0b1e2a1 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -1,4 +1,9 @@ -"""Support for devices connected to UniFi POE.""" +"""Switch platform for UniFi integration. + +Support for controlling power supply of clients which are powered over Ethernet (POE). +Support for controlling network access of clients selected in option flow. +Support for controlling deep packet inspection (DPI) restriction groups. +""" import logging from typing import Any diff --git a/tests/components/unifi/test_config_flow.py b/tests/components/unifi/test_config_flow.py index 15220e68914f..790e204c1dd4 100644 --- a/tests/components/unifi/test_config_flow.py +++ b/tests/components/unifi/test_config_flow.py @@ -338,32 +338,6 @@ async def test_flow_fails_controller_unavailable(hass, aioclient_mock): assert result["errors"] == {"base": "service_unavailable"} -async def test_flow_fails_unknown_problem(hass, aioclient_mock): - """Test config flow.""" - result = await hass.config_entries.flow.async_init( - UNIFI_DOMAIN, context={"source": "user"} - ) - - assert result["type"] == data_entry_flow.RESULT_TYPE_FORM - assert result["step_id"] == "user" - - aioclient_mock.get("https://1.2.3.4:1234", status=302) - - with patch("aiounifi.Controller.login", side_effect=Exception): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], - user_input={ - CONF_HOST: "1.2.3.4", - CONF_USERNAME: "username", - CONF_PASSWORD: "password", - CONF_PORT: 1234, - CONF_VERIFY_SSL: True, - }, - ) - - assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT - - async def test_reauth_flow_update_configuration(hass, aioclient_mock): """Verify reauth flow can update controller configuration.""" controller = await setup_unifi_integration(hass) diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index 6acd507eaad7..e484e041a883 100644 --- a/tests/components/unifi/test_controller.py +++ b/tests/components/unifi/test_controller.py @@ -201,9 +201,11 @@ async def test_controller_setup(hass): assert controller.mac is None - assert controller.signal_update == "unifi-update-1.2.3.4-site_id" - assert controller.signal_remove == "unifi-remove-1.2.3.4-site_id" - assert controller.signal_options_update == "unifi-options-1.2.3.4-site_id" + assert controller.signal_reachable == "unifi-reachable-1" + assert controller.signal_update == "unifi-update-1" + assert controller.signal_remove == "unifi-remove-1" + assert controller.signal_options_update == "unifi-options-1" + assert controller.signal_heartbeat_missed == "unifi-heartbeat-missed" async def test_controller_mac(hass):