UniFi refactor using site data (#98549)

* Clean up

* Simplify admin verification

* Streamline using sites in config_flow

* Bump aiounifi

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(

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:
+        if len(self.sites.values()) == 1:
-            return await self.async_step_site(
+            return await self.async_step_site({CONF_SITE_ID: next(iter(self.sites))})
-                {CONF_SITE_ID: next(iter(self.site_names))}
-            )
 
+        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(
+            data_schema=vol.Schema({vol.Required(CONF_SITE_ID): vol.In(site_names)}),
-                {vol.Required(CONF_SITE_ID): vol.In(self.site_names)}
-            ),
         )
 
     async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult:

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 = config_entry.data[CONF_SITE_ID]
-        self._site_name: str | None = None
+        self.is_admin = False
-        self._site_role: str | None = None
 
         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()
+        assert self.config_entry.unique_id is not None
-        for site in sites.values():
+        self.is_admin = self.api.sites[self.config_entry.unique_id].role == "admin"
-            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"]
 
         # Restore clients that are not a part of active clients list.
         entity_registry = er.async_get(self.hass)

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

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(

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

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:

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)

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

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

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.is_admin == (SITE[0]["role"] == "admin")
-    assert controller.site_name == SITE[0]["desc"]
-    assert controller.site_role == SITE[0]["role"]
 
     assert controller.option_allow_bandwidth_sensors == DEFAULT_ALLOW_BANDWIDTH_SENSORS
     assert controller.option_allow_uptime_sensors == DEFAULT_ALLOW_UPTIME_SENSORS

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,

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]
 
 

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)
+    site = deepcopy(SITE)
-    description[0]["site_role"] = "not admin"
+    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.call_count == 12
-    assert aioclient_mock.mock_calls[10][2] == {
+    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.call_count == 13
-    assert aioclient_mock.mock_calls[11][2] == {
+    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.call_count == 14
-    assert aioclient_mock.mock_calls[12][2] == {"enabled": False}
+    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.call_count == 15
-    assert aioclient_mock.mock_calls[13][2] == {"enabled": True}
+    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.call_count == 12
-    assert aioclient_mock.mock_calls[10][2] == {
+    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.call_count == 13
-    assert aioclient_mock.mock_calls[11][2] == {
+    assert aioclient_mock.mock_calls[12][2] == {
         "mac": "00:00:00:00:01:01",
         "cmd": "unblock-sta",
     }

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)
+    site = deepcopy(SITE)
-    description[0]["site_role"] = "not admin"
+    site[0]["role"] = "not admin"
 
     await setup_unifi_integration(
-        hass,
+        hass, aioclient_mock, sites=site, devices_response=[DEVICE_1]
-        aioclient_mock,
-        site_description=description,
-        devices_response=[DEVICE_1],
     )
 
     assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 1