diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index a9d39251838..5f19a01ce45 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/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) diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index 07b81621750..1d89215dc89 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/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( - self.reauth_config_entry, data=data - ) - 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 + config_entry = self.reauth_config_entry + abort_reason = "reauth_successful" + 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: - return await self.async_step_site({CONF_SITE_ID: next(iter(self.sites))}) + if len(self.site_names) == 1: + 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, ) diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index bd55f4479fa..5d5e679e75e 100644 --- a/homeassistant/components/unifi/controller.py +++ b/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 diff --git a/tests/components/unifi/test_config_flow.py b/tests/components/unifi/test_config_flow.py index 302564dbbd5..096e6ba7791 100644 --- a/tests/components/unifi/test_config_flow.py +++ b/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": "site2", "role": "admin", "desc": "site2 name"}, + {"name": "default", "role": "admin", "desc": "site name", "_id": "1"}, + {"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": "site2"}) + assert result["data_schema"]({"site": "1"}) + 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}, diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index 077481a9320..3ecd44b3db7 100644 --- a/tests/components/unifi/test_controller.py +++ b/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) diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py index 841e9ec7576..9de8b0a0990 100644 --- a/tests/components/unifi/test_init.py +++ b/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()