UniFi refactor using site data (#98549)
* Clean up * Simplify admin verification * Streamline using sites in config_flow * Bump aiounifi
This commit is contained in:
parent
7827f9ccae
commit
9e42451934
15 changed files with 65 additions and 91 deletions
|
@ -89,7 +89,7 @@ async def async_setup_entry(
|
|||
"""Set up button platform for UniFi Network integration."""
|
||||
controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
|
||||
|
||||
if controller.site_role != "admin":
|
||||
if not controller.is_admin:
|
||||
return
|
||||
|
||||
controller.register_platform_add_entities(
|
||||
|
|
|
@ -13,6 +13,7 @@ from types import MappingProxyType
|
|||
from typing import Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from aiounifi.interfaces.sites import Sites
|
||||
import voluptuous as vol
|
||||
|
||||
from homeassistant import config_entries
|
||||
|
@ -63,6 +64,8 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
|
|||
|
||||
VERSION = 1
|
||||
|
||||
sites: Sites
|
||||
|
||||
@staticmethod
|
||||
@callback
|
||||
def async_get_options_flow(
|
||||
|
@ -74,8 +77,6 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
|
|||
def __init__(self) -> None:
|
||||
"""Initialize the UniFi Network flow."""
|
||||
self.config: dict[str, Any] = {}
|
||||
self.site_ids: dict[str, str] = {}
|
||||
self.site_names: dict[str, str] = {}
|
||||
self.reauth_config_entry: config_entries.ConfigEntry | None = None
|
||||
self.reauth_schema: dict[vol.Marker, Any] = {}
|
||||
|
||||
|
@ -99,7 +100,8 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
|
|||
controller = await get_unifi_controller(
|
||||
self.hass, MappingProxyType(self.config)
|
||||
)
|
||||
sites = await controller.sites()
|
||||
await controller.sites.update()
|
||||
self.sites = controller.sites
|
||||
|
||||
except AuthenticationRequired:
|
||||
errors["base"] = "faulty_credentials"
|
||||
|
@ -108,12 +110,10 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
|
|||
errors["base"] = "service_unavailable"
|
||||
|
||||
else:
|
||||
self.site_ids = {site["_id"]: site["name"] for site in sites.values()}
|
||||
self.site_names = {site["_id"]: site["desc"] for site in sites.values()}
|
||||
|
||||
if (
|
||||
self.reauth_config_entry
|
||||
and self.reauth_config_entry.unique_id in self.site_names
|
||||
and self.reauth_config_entry.unique_id is not None
|
||||
and self.reauth_config_entry.unique_id in self.sites
|
||||
):
|
||||
return await self.async_step_site(
|
||||
{CONF_SITE_ID: self.reauth_config_entry.unique_id}
|
||||
|
@ -148,7 +148,7 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
|
|||
"""Select site to control."""
|
||||
if user_input is not None:
|
||||
unique_id = user_input[CONF_SITE_ID]
|
||||
self.config[CONF_SITE_ID] = self.site_ids[unique_id]
|
||||
self.config[CONF_SITE_ID] = self.sites[unique_id].name
|
||||
|
||||
config_entry = await self.async_set_unique_id(unique_id)
|
||||
abort_reason = "configuration_updated"
|
||||
|
@ -171,19 +171,16 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
|
|||
await self.hass.config_entries.async_reload(config_entry.entry_id)
|
||||
return self.async_abort(reason=abort_reason)
|
||||
|
||||
site_nice_name = self.site_names[unique_id]
|
||||
site_nice_name = self.sites[unique_id].description
|
||||
return self.async_create_entry(title=site_nice_name, data=self.config)
|
||||
|
||||
if len(self.site_names) == 1:
|
||||
return await self.async_step_site(
|
||||
{CONF_SITE_ID: next(iter(self.site_names))}
|
||||
)
|
||||
if len(self.sites.values()) == 1:
|
||||
return await self.async_step_site({CONF_SITE_ID: next(iter(self.sites))})
|
||||
|
||||
site_names = {site.site_id: site.description for site in self.sites.values()}
|
||||
return self.async_show_form(
|
||||
step_id="site",
|
||||
data_schema=vol.Schema(
|
||||
{vol.Required(CONF_SITE_ID): vol.In(self.site_names)}
|
||||
),
|
||||
data_schema=vol.Schema({vol.Required(CONF_SITE_ID): vol.In(site_names)}),
|
||||
)
|
||||
|
||||
async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult:
|
||||
|
|
|
@ -87,9 +87,8 @@ class UniFiController:
|
|||
self.available = True
|
||||
self.wireless_clients = hass.data[UNIFI_WIRELESS_CLIENTS]
|
||||
|
||||
self.site_id: str = ""
|
||||
self._site_name: str | None = None
|
||||
self._site_role: str | None = None
|
||||
self.site = config_entry.data[CONF_SITE_ID]
|
||||
self.is_admin = False
|
||||
|
||||
self._cancel_heartbeat_check: CALLBACK_TYPE | None = None
|
||||
self._heartbeat_time: dict[str, datetime] = {}
|
||||
|
@ -154,22 +153,6 @@ class UniFiController:
|
|||
host: str = self.config_entry.data[CONF_HOST]
|
||||
return host
|
||||
|
||||
@property
|
||||
def site(self) -> str:
|
||||
"""Return the site of this config entry."""
|
||||
site_id: str = self.config_entry.data[CONF_SITE_ID]
|
||||
return site_id
|
||||
|
||||
@property
|
||||
def site_name(self) -> str | None:
|
||||
"""Return the nice name of site."""
|
||||
return self._site_name
|
||||
|
||||
@property
|
||||
def site_role(self) -> str | None:
|
||||
"""Return the site user role of this controller."""
|
||||
return self._site_role
|
||||
|
||||
@property
|
||||
def mac(self) -> str | None:
|
||||
"""Return the mac address of this controller."""
|
||||
|
@ -264,15 +247,8 @@ class UniFiController:
|
|||
"""Set up a UniFi Network instance."""
|
||||
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"]
|
||||
assert self.config_entry.unique_id is not None
|
||||
self.is_admin = self.api.sites[self.config_entry.unique_id].role == "admin"
|
||||
|
||||
# Restore clients that are not a part of active clients list.
|
||||
entity_registry = er.async_get(self.hass)
|
||||
|
|
|
@ -94,7 +94,7 @@ async def async_get_config_entry_diagnostics(
|
|||
diag["config"] = async_redact_data(
|
||||
async_replace_dict_data(config_entry.as_dict(), macs_to_redact), REDACT_CONFIG
|
||||
)
|
||||
diag["site_role"] = controller.site_role
|
||||
diag["role_is_admin"] = controller.is_admin
|
||||
diag["clients"] = {
|
||||
macs_to_redact[k]: async_redact_data(
|
||||
async_replace_dict_data(v.raw, macs_to_redact), REDACT_CLIENTS
|
||||
|
|
|
@ -85,7 +85,7 @@ async def async_setup_entry(
|
|||
"""Set up image platform for UniFi Network integration."""
|
||||
controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
|
||||
|
||||
if controller.site_role != "admin":
|
||||
if not controller.is_admin:
|
||||
return
|
||||
|
||||
controller.register_platform_add_entities(
|
||||
|
|
|
@ -8,7 +8,7 @@
|
|||
"iot_class": "local_push",
|
||||
"loggers": ["aiounifi"],
|
||||
"quality_scale": "platinum",
|
||||
"requirements": ["aiounifi==53"],
|
||||
"requirements": ["aiounifi==55"],
|
||||
"ssdp": [
|
||||
{
|
||||
"manufacturer": "Ubiquiti Networks",
|
||||
|
|
|
@ -274,7 +274,7 @@ async def async_setup_entry(
|
|||
"""Set up switches for UniFi Network integration."""
|
||||
controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
|
||||
|
||||
if controller.site_role != "admin":
|
||||
if not controller.is_admin:
|
||||
return
|
||||
|
||||
for mac in controller.option_block_clients:
|
||||
|
|
|
@ -103,7 +103,7 @@ class UnifiDeviceUpdateEntity(UnifiEntity[_HandlerT, _DataT], UpdateEntity):
|
|||
def async_initiate_state(self) -> None:
|
||||
"""Initiate entity state."""
|
||||
self._attr_supported_features = UpdateEntityFeature.PROGRESS
|
||||
if self.controller.site_role == "admin":
|
||||
if self.controller.is_admin:
|
||||
self._attr_supported_features |= UpdateEntityFeature.INSTALL
|
||||
|
||||
self.async_update_state(ItemEvent.ADDED, self._obj_id)
|
||||
|
|
|
@ -360,7 +360,7 @@ aiosyncthing==0.5.1
|
|||
aiotractive==0.5.5
|
||||
|
||||
# homeassistant.components.unifi
|
||||
aiounifi==53
|
||||
aiounifi==55
|
||||
|
||||
# homeassistant.components.vlc_telnet
|
||||
aiovlc==0.1.0
|
||||
|
|
|
@ -335,7 +335,7 @@ aiosyncthing==0.5.1
|
|||
aiotractive==0.5.5
|
||||
|
||||
# homeassistant.components.unifi
|
||||
aiounifi==53
|
||||
aiounifi==55
|
||||
|
||||
# homeassistant.components.vlc_telnet
|
||||
aiovlc==0.1.0
|
||||
|
|
|
@ -80,7 +80,6 @@ ENTRY_OPTIONS = {}
|
|||
CONFIGURATION = []
|
||||
|
||||
SITE = [{"desc": "Site name", "name": "site_id", "role": "admin", "_id": "1"}]
|
||||
DESCRIPTION = [{"name": "username", "site_name": "site_id", "site_role": "admin"}]
|
||||
|
||||
|
||||
def mock_default_unifi_requests(
|
||||
|
@ -88,12 +87,13 @@ def mock_default_unifi_requests(
|
|||
host,
|
||||
site_id,
|
||||
sites=None,
|
||||
description=None,
|
||||
clients_response=None,
|
||||
clients_all_response=None,
|
||||
devices_response=None,
|
||||
dpiapp_response=None,
|
||||
dpigroup_response=None,
|
||||
port_forward_response=None,
|
||||
system_information_response=None,
|
||||
wlans_response=None,
|
||||
):
|
||||
"""Mock default UniFi requests responses."""
|
||||
|
@ -111,12 +111,6 @@ def mock_default_unifi_requests(
|
|||
headers={"content-type": CONTENT_TYPE_JSON},
|
||||
)
|
||||
|
||||
aioclient_mock.get(
|
||||
f"https://{host}:1234/api/s/{site_id}/self",
|
||||
json={"data": description or [], "meta": {"rc": "ok"}},
|
||||
headers={"content-type": CONTENT_TYPE_JSON},
|
||||
)
|
||||
|
||||
aioclient_mock.get(
|
||||
f"https://{host}:1234/api/s/{site_id}/stat/sta",
|
||||
json={"data": clients_response or [], "meta": {"rc": "ok"}},
|
||||
|
@ -142,6 +136,16 @@ def mock_default_unifi_requests(
|
|||
json={"data": dpigroup_response or [], "meta": {"rc": "ok"}},
|
||||
headers={"content-type": CONTENT_TYPE_JSON},
|
||||
)
|
||||
aioclient_mock.get(
|
||||
f"https://{host}:1234/api/s/{site_id}/rest/portforward",
|
||||
json={"data": port_forward_response or [], "meta": {"rc": "ok"}},
|
||||
headers={"content-type": CONTENT_TYPE_JSON},
|
||||
)
|
||||
aioclient_mock.get(
|
||||
f"https://{host}:1234/api/s/{site_id}/stat/sysinfo",
|
||||
json={"data": system_information_response or [], "meta": {"rc": "ok"}},
|
||||
headers={"content-type": CONTENT_TYPE_JSON},
|
||||
)
|
||||
aioclient_mock.get(
|
||||
f"https://{host}:1234/api/s/{site_id}/rest/wlanconf",
|
||||
json={"data": wlans_response or [], "meta": {"rc": "ok"}},
|
||||
|
@ -156,12 +160,13 @@ async def setup_unifi_integration(
|
|||
config=ENTRY_CONFIG,
|
||||
options=ENTRY_OPTIONS,
|
||||
sites=SITE,
|
||||
site_description=DESCRIPTION,
|
||||
clients_response=None,
|
||||
clients_all_response=None,
|
||||
devices_response=None,
|
||||
dpiapp_response=None,
|
||||
dpigroup_response=None,
|
||||
port_forward_response=None,
|
||||
system_information_response=None,
|
||||
wlans_response=None,
|
||||
known_wireless_clients=None,
|
||||
controllers=None,
|
||||
|
@ -192,12 +197,13 @@ async def setup_unifi_integration(
|
|||
host=config_entry.data[CONF_HOST],
|
||||
site_id=config_entry.data[CONF_SITE_ID],
|
||||
sites=sites,
|
||||
description=site_description,
|
||||
clients_response=clients_response,
|
||||
clients_all_response=clients_all_response,
|
||||
devices_response=devices_response,
|
||||
dpiapp_response=dpiapp_response,
|
||||
dpigroup_response=dpigroup_response,
|
||||
port_forward_response=port_forward_response,
|
||||
system_information_response=system_information_response,
|
||||
wlans_response=wlans_response,
|
||||
)
|
||||
|
||||
|
@ -230,9 +236,7 @@ async def test_controller_setup(
|
|||
assert forward_entry_setup.mock_calls[4][1] == (entry, SWITCH_DOMAIN)
|
||||
|
||||
assert controller.host == ENTRY_CONFIG[CONF_HOST]
|
||||
assert controller.site == ENTRY_CONFIG[CONF_SITE_ID]
|
||||
assert controller.site_name == SITE[0]["desc"]
|
||||
assert controller.site_role == SITE[0]["role"]
|
||||
assert controller.is_admin == (SITE[0]["role"] == "admin")
|
||||
|
||||
assert controller.option_allow_bandwidth_sensors == DEFAULT_ALLOW_BANDWIDTH_SENSORS
|
||||
assert controller.option_allow_uptime_sensors == DEFAULT_ALLOW_UPTIME_SENSORS
|
||||
|
|
|
@ -141,7 +141,7 @@ async def test_entry_diagnostics(
|
|||
"unique_id": "1",
|
||||
"version": 1,
|
||||
},
|
||||
"site_role": "admin",
|
||||
"role_is_admin": True,
|
||||
"clients": {
|
||||
"00:00:00:00:00:00": {
|
||||
"blocked": False,
|
||||
|
|
|
@ -24,7 +24,7 @@ async def test_successful_config_entry(
|
|||
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
|
||||
) -> None:
|
||||
"""Test that configured options for a host are loaded via config entry."""
|
||||
await setup_unifi_integration(hass, aioclient_mock, unique_id=None)
|
||||
await setup_unifi_integration(hass, aioclient_mock)
|
||||
assert hass.data[UNIFI_DOMAIN]
|
||||
|
||||
|
||||
|
|
|
@ -36,8 +36,8 @@ from homeassistant.util import dt as dt_util
|
|||
|
||||
from .test_controller import (
|
||||
CONTROLLER_HOST,
|
||||
DESCRIPTION,
|
||||
ENTRY_CONFIG,
|
||||
SITE,
|
||||
setup_unifi_integration,
|
||||
)
|
||||
|
||||
|
@ -778,7 +778,7 @@ async def test_no_clients(
|
|||
},
|
||||
)
|
||||
|
||||
assert aioclient_mock.call_count == 10
|
||||
assert aioclient_mock.call_count == 11
|
||||
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
|
||||
|
||||
|
||||
|
@ -803,13 +803,13 @@ async def test_not_admin(
|
|||
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
|
||||
) -> None:
|
||||
"""Test that switch platform only work on an admin account."""
|
||||
description = deepcopy(DESCRIPTION)
|
||||
description[0]["site_role"] = "not admin"
|
||||
site = deepcopy(SITE)
|
||||
site[0]["role"] = "not admin"
|
||||
await setup_unifi_integration(
|
||||
hass,
|
||||
aioclient_mock,
|
||||
options={CONF_TRACK_CLIENTS: False, CONF_TRACK_DEVICES: False},
|
||||
site_description=description,
|
||||
sites=site,
|
||||
clients_response=[CLIENT_1],
|
||||
devices_response=[DEVICE_1],
|
||||
)
|
||||
|
@ -867,8 +867,8 @@ async def test_switches(
|
|||
await hass.services.async_call(
|
||||
SWITCH_DOMAIN, "turn_off", {"entity_id": "switch.block_client_1"}, blocking=True
|
||||
)
|
||||
assert aioclient_mock.call_count == 11
|
||||
assert aioclient_mock.mock_calls[10][2] == {
|
||||
assert aioclient_mock.call_count == 12
|
||||
assert aioclient_mock.mock_calls[11][2] == {
|
||||
"mac": "00:00:00:00:01:01",
|
||||
"cmd": "block-sta",
|
||||
}
|
||||
|
@ -876,8 +876,8 @@ async def test_switches(
|
|||
await hass.services.async_call(
|
||||
SWITCH_DOMAIN, "turn_on", {"entity_id": "switch.block_client_1"}, blocking=True
|
||||
)
|
||||
assert aioclient_mock.call_count == 12
|
||||
assert aioclient_mock.mock_calls[11][2] == {
|
||||
assert aioclient_mock.call_count == 13
|
||||
assert aioclient_mock.mock_calls[12][2] == {
|
||||
"mac": "00:00:00:00:01:01",
|
||||
"cmd": "unblock-sta",
|
||||
}
|
||||
|
@ -894,8 +894,8 @@ async def test_switches(
|
|||
{"entity_id": "switch.block_media_streaming"},
|
||||
blocking=True,
|
||||
)
|
||||
assert aioclient_mock.call_count == 13
|
||||
assert aioclient_mock.mock_calls[12][2] == {"enabled": False}
|
||||
assert aioclient_mock.call_count == 14
|
||||
assert aioclient_mock.mock_calls[13][2] == {"enabled": False}
|
||||
|
||||
await hass.services.async_call(
|
||||
SWITCH_DOMAIN,
|
||||
|
@ -903,8 +903,8 @@ async def test_switches(
|
|||
{"entity_id": "switch.block_media_streaming"},
|
||||
blocking=True,
|
||||
)
|
||||
assert aioclient_mock.call_count == 14
|
||||
assert aioclient_mock.mock_calls[13][2] == {"enabled": True}
|
||||
assert aioclient_mock.call_count == 15
|
||||
assert aioclient_mock.mock_calls[14][2] == {"enabled": True}
|
||||
|
||||
|
||||
async def test_remove_switches(
|
||||
|
@ -990,8 +990,8 @@ async def test_block_switches(
|
|||
await hass.services.async_call(
|
||||
SWITCH_DOMAIN, "turn_off", {"entity_id": "switch.block_client_1"}, blocking=True
|
||||
)
|
||||
assert aioclient_mock.call_count == 11
|
||||
assert aioclient_mock.mock_calls[10][2] == {
|
||||
assert aioclient_mock.call_count == 12
|
||||
assert aioclient_mock.mock_calls[11][2] == {
|
||||
"mac": "00:00:00:00:01:01",
|
||||
"cmd": "block-sta",
|
||||
}
|
||||
|
@ -999,8 +999,8 @@ async def test_block_switches(
|
|||
await hass.services.async_call(
|
||||
SWITCH_DOMAIN, "turn_on", {"entity_id": "switch.block_client_1"}, blocking=True
|
||||
)
|
||||
assert aioclient_mock.call_count == 12
|
||||
assert aioclient_mock.mock_calls[11][2] == {
|
||||
assert aioclient_mock.call_count == 13
|
||||
assert aioclient_mock.mock_calls[12][2] == {
|
||||
"mac": "00:00:00:00:01:01",
|
||||
"cmd": "unblock-sta",
|
||||
}
|
||||
|
|
|
@ -26,7 +26,7 @@ from homeassistant.const import (
|
|||
)
|
||||
from homeassistant.core import HomeAssistant
|
||||
|
||||
from .test_controller import DESCRIPTION, setup_unifi_integration
|
||||
from .test_controller import SITE, setup_unifi_integration
|
||||
|
||||
from tests.test_util.aiohttp import AiohttpClientMocker
|
||||
|
||||
|
@ -136,14 +136,11 @@ async def test_not_admin(
|
|||
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
|
||||
) -> None:
|
||||
"""Test that the INSTALL feature is not available on a non-admin account."""
|
||||
description = deepcopy(DESCRIPTION)
|
||||
description[0]["site_role"] = "not admin"
|
||||
site = deepcopy(SITE)
|
||||
site[0]["role"] = "not admin"
|
||||
|
||||
await setup_unifi_integration(
|
||||
hass,
|
||||
aioclient_mock,
|
||||
site_description=description,
|
||||
devices_response=[DEVICE_1],
|
||||
hass, aioclient_mock, sites=site, devices_response=[DEVICE_1]
|
||||
)
|
||||
|
||||
assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 1
|
||||
|
|
Loading…
Add table
Reference in a new issue