Add unique id to UniFi config entries using the unique id of the site it is controlling (#45737)

* Add unique id to UniFi config entries using the unique id of the site it is controlling

* Fix failing test

homeassistant/components/unifi/__init__.py

@@ -32,6 +32,11 @@ async def async_setup_entry(hass, config_entry):
     if not await controller.async_setup():
         return False
 
+    if config_entry.unique_id is None:
+        hass.config_entries.async_update_entry(
+            config_entry, unique_id=controller.site_id
+        )
+
     hass.data[UNIFI_DOMAIN][config_entry.entry_id] = controller
 
     hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, controller.shutdown)

homeassistant/components/unifi/config_flow.py

@@ -70,7 +70,8 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
     def __init__(self):
         """Initialize the UniFi flow."""
         self.config = {}
-        self.sites = None
+        self.site_ids = {}
+        self.site_names = {}
         self.reauth_config_entry = None
         self.reauth_config = {}
         self.reauth_schema = {}
@@ -101,11 +102,15 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
                 errors["base"] = "service_unavailable"
 
             else:
-                self.sites = {site["name"]: site["desc"] for site in sites.values()}
+                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.get(CONF_SITE_ID) in self.sites:
+                if (
+                    self.reauth_config_entry
+                    and self.reauth_config_entry.unique_id in self.site_names
+                ):
                     return await self.async_step_site(
-                        {CONF_SITE_ID: self.reauth_config[CONF_SITE_ID]}
+                        {CONF_SITE_ID: self.reauth_config_entry.unique_id}
                     )
 
                 return await self.async_step_site()
@@ -136,26 +141,18 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
 
         if user_input is not None:
 
-            self.config[CONF_SITE_ID] = user_input[CONF_SITE_ID]
+            unique_id = user_input[CONF_SITE_ID]
+            self.config[CONF_SITE_ID] = self.site_ids[unique_id]
             data = {CONF_CONTROLLER: self.config}
 
+            config_entry = await self.async_set_unique_id(unique_id)
+            abort_reason = "configuration_updated"
+
             if self.reauth_config_entry:
-                self.hass.config_entries.async_update_entry(
+                config_entry = self.reauth_config_entry
-                    self.reauth_config_entry, data=data
+                abort_reason = "reauth_successful"
-                )
-                await self.hass.config_entries.async_reload(
-                    self.reauth_config_entry.entry_id
-                )
-                return self.async_abort(reason="reauth_successful")
-
-            for config_entry in self._async_current_entries():
-                controller_data = config_entry.data[CONF_CONTROLLER]
-                if (
-                    controller_data[CONF_HOST] != self.config[CONF_HOST]
-                    or controller_data[CONF_SITE_ID] != self.config[CONF_SITE_ID]
-                ):
-                    continue
 
+            if config_entry:
                 controller = self.hass.data.get(UNIFI_DOMAIN, {}).get(
                     config_entry.entry_id
                 )
@@ -165,17 +162,21 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
 
                 self.hass.config_entries.async_update_entry(config_entry, data=data)
                 await self.hass.config_entries.async_reload(config_entry.entry_id)
-                return self.async_abort(reason="configuration_updated")
+                return self.async_abort(reason=abort_reason)
 
-            site_nice_name = self.sites[self.config[CONF_SITE_ID]]
+            site_nice_name = self.site_names[unique_id]
             return self.async_create_entry(title=site_nice_name, data=data)
 
-        if len(self.sites) == 1:
+        if len(self.site_names) == 1:
-            return await self.async_step_site({CONF_SITE_ID: next(iter(self.sites))})
+            return await self.async_step_site(
+                {CONF_SITE_ID: next(iter(self.site_names))}
+            )
 
         return self.async_show_form(
             step_id="site",
-            data_schema=vol.Schema({vol.Required(CONF_SITE_ID): vol.In(self.sites)}),
+            data_schema=vol.Schema(
+                {vol.Required(CONF_SITE_ID): vol.In(self.site_names)}
+            ),
             errors=errors,
         )
 

homeassistant/components/unifi/controller.py

@@ -95,6 +95,7 @@ class UniFiController:
         self.wireless_clients = None
 
         self.listeners = []
+        self.site_id: str = ""
         self._site_name = None
         self._site_role = None
 
@@ -321,6 +322,7 @@ class UniFiController:
 
         for site in sites.values():
             if self.site == site["name"]:
+                self.site_id = site["_id"]
                 self._site_name = site["desc"]
                 break
 

tests/components/unifi/test_config_flow.py

@@ -112,7 +112,9 @@ async def test_flow_works(hass, aioclient_mock, mock_discovery):
     aioclient_mock.get(
         "https://1.2.3.4:1234/api/self/sites",
         json={
-            "data": [{"desc": "Site name", "name": "site_id", "role": "admin"}],
+            "data": [
+                {"desc": "Site name", "name": "site_id", "role": "admin", "_id": "1"}
+            ],
             "meta": {"rc": "ok"},
         },
         headers={"content-type": CONTENT_TYPE_JSON},
@@ -143,7 +145,7 @@ async def test_flow_works(hass, aioclient_mock, mock_discovery):
     }
 
 
-async def test_flow_works_multiple_sites(hass, aioclient_mock):
+async def test_flow_multiple_sites(hass, aioclient_mock):
     """Test config flow works when finding multiple sites."""
     result = await hass.config_entries.flow.async_init(
         UNIFI_DOMAIN, context={"source": "user"}
@@ -164,8 +166,8 @@ async def test_flow_works_multiple_sites(hass, aioclient_mock):
         "https://1.2.3.4:1234/api/self/sites",
         json={
             "data": [
-                {"name": "default", "role": "admin", "desc": "site name"},
+                {"name": "default", "role": "admin", "desc": "site name", "_id": "1"},
-                {"name": "site2", "role": "admin", "desc": "site2 name"},
+                {"name": "site2", "role": "admin", "desc": "site2 name", "_id": "2"},
             ],
             "meta": {"rc": "ok"},
         },
@@ -185,8 +187,8 @@ async def test_flow_works_multiple_sites(hass, aioclient_mock):
 
     assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
     assert result["step_id"] == "site"
-    assert result["data_schema"]({"site": "default"})
+    assert result["data_schema"]({"site": "1"})
-    assert result["data_schema"]({"site": "site2"})
+    assert result["data_schema"]({"site": "2"})
 
 
 async def test_flow_raise_already_configured(hass, aioclient_mock):
@@ -213,7 +215,9 @@ async def test_flow_raise_already_configured(hass, aioclient_mock):
     aioclient_mock.get(
         "https://1.2.3.4:1234/api/self/sites",
         json={
-            "data": [{"desc": "Site name", "name": "site_id", "role": "admin"}],
+            "data": [
+                {"desc": "Site name", "name": "site_id", "role": "admin", "_id": "1"}
+            ],
             "meta": {"rc": "ok"},
         },
         headers={"content-type": CONTENT_TYPE_JSON},
@@ -237,12 +241,16 @@ async def test_flow_raise_already_configured(hass, aioclient_mock):
 async def test_flow_aborts_configuration_updated(hass, aioclient_mock):
     """Test config flow aborts since a connected config entry already exists."""
     entry = MockConfigEntry(
-        domain=UNIFI_DOMAIN, data={"controller": {"host": "1.2.3.4", "site": "office"}}
+        domain=UNIFI_DOMAIN,
+        data={"controller": {"host": "1.2.3.4", "site": "office"}},
+        unique_id="2",
     )
     entry.add_to_hass(hass)
 
     entry = MockConfigEntry(
-        domain=UNIFI_DOMAIN, data={"controller": {"host": "1.2.3.4", "site": "site_id"}}
+        domain=UNIFI_DOMAIN,
+        data={"controller": {"host": "1.2.3.4", "site": "site_id"}},
+        unique_id="1",
     )
     entry.add_to_hass(hass)
 
@@ -264,7 +272,9 @@ async def test_flow_aborts_configuration_updated(hass, aioclient_mock):
     aioclient_mock.get(
         "https://1.2.3.4:1234/api/self/sites",
         json={
-            "data": [{"desc": "Site name", "name": "site_id", "role": "admin"}],
+            "data": [
+                {"desc": "Site name", "name": "site_id", "role": "admin", "_id": "1"}
+            ],
             "meta": {"rc": "ok"},
         },
         headers={"content-type": CONTENT_TYPE_JSON},
@@ -343,6 +353,8 @@ async def test_flow_fails_controller_unavailable(hass, aioclient_mock):
 async def test_reauth_flow_update_configuration(hass, aioclient_mock):
     """Verify reauth flow can update controller configuration."""
     config_entry = await setup_unifi_integration(hass, aioclient_mock)
+    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+    controller.available = False
 
     result = await hass.config_entries.flow.async_init(
         UNIFI_DOMAIN,
@@ -366,7 +378,9 @@ async def test_reauth_flow_update_configuration(hass, aioclient_mock):
     aioclient_mock.get(
         "https://1.2.3.4:1234/api/self/sites",
         json={
-            "data": [{"desc": "Site name", "name": "site_id", "role": "admin"}],
+            "data": [
+                {"desc": "Site name", "name": "site_id", "role": "admin", "_id": "1"}
+            ],
             "meta": {"rc": "ok"},
         },
         headers={"content-type": CONTENT_TYPE_JSON},

tests/components/unifi/test_controller.py

@@ -71,7 +71,7 @@ ENTRY_OPTIONS = {}
 
 CONFIGURATION = []
 
-SITE = [{"desc": "Site name", "name": "site_id", "role": "admin"}]
+SITE = [{"desc": "Site name", "name": "site_id", "role": "admin", "_id": "1"}]
 DESCRIPTION = [{"name": "username", "site_name": "site_id", "site_role": "admin"}]
 
 
@@ -166,6 +166,7 @@ async def setup_unifi_integration(
         data=deepcopy(config),
         options=deepcopy(options),
         entry_id=1,
+        unique_id="1",
     )
     config_entry.add_to_hass(hass)
 

tests/components/unifi/test_init.py

@@ -45,6 +45,7 @@ async def test_controller_no_mac(hass):
                 "verify_ssl": True,
             },
         },
+        unique_id="1",
     )
     entry.add_to_hass(hass)
     mock_registry = Mock()