From 9e42451934d0513eb53fbe84dfe5576f7727d695 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Fri, 18 Aug 2023 22:44:59 +0200 Subject: [PATCH] UniFi refactor using site data (#98549) * Clean up * Simplify admin verification * Streamline using sites in config_flow * Bump aiounifi --- homeassistant/components/unifi/button.py | 2 +- homeassistant/components/unifi/config_flow.py | 29 +++++++--------- homeassistant/components/unifi/controller.py | 32 +++-------------- homeassistant/components/unifi/diagnostics.py | 2 +- homeassistant/components/unifi/image.py | 2 +- homeassistant/components/unifi/manifest.json | 2 +- homeassistant/components/unifi/switch.py | 2 +- homeassistant/components/unifi/update.py | 2 +- requirements_all.txt | 2 +- requirements_test_all.txt | 2 +- tests/components/unifi/test_controller.py | 30 +++++++++------- tests/components/unifi/test_diagnostics.py | 2 +- tests/components/unifi/test_init.py | 2 +- tests/components/unifi/test_switch.py | 34 +++++++++---------- tests/components/unifi/test_update.py | 11 +++--- 15 files changed, 65 insertions(+), 91 deletions(-) diff --git a/homeassistant/components/unifi/button.py b/homeassistant/components/unifi/button.py index 6b0660325f0..0235f6156cc 100644 --- a/homeassistant/components/unifi/button.py +++ b/homeassistant/components/unifi/button.py @@ -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( diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index 12f2d49e416..8c0696463c5 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/homeassistant/components/unifi/config_flow.py @@ -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: diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 649d7c30fdb..c1ffa0aa57d 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -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) diff --git a/homeassistant/components/unifi/diagnostics.py b/homeassistant/components/unifi/diagnostics.py index 3c72c06d6f2..c01dc193078 100644 --- a/homeassistant/components/unifi/diagnostics.py +++ b/homeassistant/components/unifi/diagnostics.py @@ -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 diff --git a/homeassistant/components/unifi/image.py b/homeassistant/components/unifi/image.py index 3ff893838c9..8231b87ee85 100644 --- a/homeassistant/components/unifi/image.py +++ b/homeassistant/components/unifi/image.py @@ -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( diff --git a/homeassistant/components/unifi/manifest.json b/homeassistant/components/unifi/manifest.json index 8f27263b288..579e64c5862 100644 --- a/homeassistant/components/unifi/manifest.json +++ b/homeassistant/components/unifi/manifest.json @@ -8,7 +8,7 @@ "iot_class": "local_push", "loggers": ["aiounifi"], "quality_scale": "platinum", - "requirements": ["aiounifi==53"], + "requirements": ["aiounifi==55"], "ssdp": [ { "manufacturer": "Ubiquiti Networks", diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index a82b9e35d45..e2b4dda3912 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -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: diff --git a/homeassistant/components/unifi/update.py b/homeassistant/components/unifi/update.py index 661a9016bdc..6526a02da83 100644 --- a/homeassistant/components/unifi/update.py +++ b/homeassistant/components/unifi/update.py @@ -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) diff --git a/requirements_all.txt b/requirements_all.txt index 1afbd67743c..aa636c42372 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -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 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index 3d0f59d7bca..aa34d430dba 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -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 diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index 2d28240a90d..a2be388af4c 100644 --- a/tests/components/unifi/test_controller.py +++ b/tests/components/unifi/test_controller.py @@ -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 diff --git a/tests/components/unifi/test_diagnostics.py b/tests/components/unifi/test_diagnostics.py index 5248836c08a..638e79ae649 100644 --- a/tests/components/unifi/test_diagnostics.py +++ b/tests/components/unifi/test_diagnostics.py @@ -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, diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py index cce26ac84cc..a1b817d67e2 100644 --- a/tests/components/unifi/test_init.py +++ b/tests/components/unifi/test_init.py @@ -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] diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index 5344ac901b7..c091fc5cc59 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -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", } diff --git a/tests/components/unifi/test_update.py b/tests/components/unifi/test_update.py index 7cf8495b9db..e59eca371d6 100644 --- a/tests/components/unifi/test_update.py +++ b/tests/components/unifi/test_update.py @@ -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