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 descriptionspull/45829/head
parent
2136b3013f
commit
83a75b02ea
|
@ -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)
|
||||
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -4,8 +4,6 @@ import logging
|
|||
LOGGER = logging.getLogger(__package__)
|
||||
DOMAIN = "unifi"
|
||||
|
||||
CONTROLLER_ID = "{host}-{site}"
|
||||
|
||||
CONF_CONTROLLER = "controller"
|
||||
CONF_SITE_ID = "site"
|
||||
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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):
|
||||
|
|
Loading…
Reference in New Issue