1
mirror of https://github.com/home-assistant/core synced 2024-08-31 05:57:13 +02:00

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
This commit is contained in:
Robert Svensson 2021-02-01 17:55:16 +01:00 committed by GitHub
parent 2136b3013f
commit 83a75b02ea
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 89 additions and 118 deletions

View File

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

View File

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

View File

@ -4,8 +4,6 @@ import logging
LOGGER = logging.getLogger(__package__)
DOMAIN = "unifi"
CONTROLLER_ID = "{host}-{site}"
CONF_CONTROLLER = "controller"
CONF_SITE_ID = "site"

View File

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

View File

@ -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

View File

@ -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

View File

@ -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

View File

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

View File

@ -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):