evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

UniFi refactor using site data (#98549)

Browse source 
* Clean up

* Simplify admin verification

* Streamline using sites in config_flow

* Bump aiounifi
This commit is contained in:
Robert Svensson 2023-08-18 22:44:59 +02:00 • committed by GitHub
parent 7827f9ccae
commit 9e42451934
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
15 changed files with 65 additions and 91 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
homeassistant/components/unifi/button.py
View file

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

29
homeassistant/components/unifi/config_flow.py
View file

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

32
homeassistant/components/unifi/controller.py
View file

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

2
homeassistant/components/unifi/diagnostics.py
View file

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

2
homeassistant/components/unifi/image.py
View file

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

2
homeassistant/components/unifi/manifest.json
View file

@ -8,7 +8,7 @@
"iot_class": "local_push",
"loggers": ["aiounifi"],
"quality_scale": "platinum",
"requirements": ["aiounifi==53"],
"requirements": ["aiounifi==55"],
"ssdp": [
{
"manufacturer": "Ubiquiti Networks",

2
homeassistant/components/unifi/switch.py
View file

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

2
homeassistant/components/unifi/update.py
View file

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

2
requirements_all.txt
View file

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

2
requirements_test_all.txt
View file

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

30
tests/components/unifi/test_controller.py
View file

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

2
tests/components/unifi/test_diagnostics.py
View file

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

2
tests/components/unifi/test_init.py
View file

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

34
tests/components/unifi/test_switch.py
View file

@ -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",
}

11
tests/components/unifi/test_update.py
View file

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