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

@@ -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)
+        data = self.data.get(config_entry.entry_id, {"wireless_devices": []})
-        key = config_entry.entry_id
-        if controller_id in self.data:
-            key = controller_id
-
-        data = self.data.get(key, {"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)
 

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()

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"
 

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:
+        unload_ok = all(
-            await self.hass.config_entries.async_forward_entry_unload(
+            await asyncio.gather(
-                self.config_entry, platform
+                *[
+                    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()

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

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

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
 

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)

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_reachable == "unifi-reachable-1"
-    assert controller.signal_remove == "unifi-remove-1.2.3.4-site_id"
+    assert controller.signal_update == "unifi-update-1"
-    assert controller.signal_options_update == "unifi-options-1.2.3.4-site_id"
+    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):