Simplify UniFi entry configuration data (#45759)

* Simplify configuration structure by removing the controller key

* Fix flake8

* Fix review comments

* Don't use migrate_entry mechanism to flatten configuration
Keep legacy configuration when creating new entries as well

homeassistant/components/unifi/__init__.py

@@ -5,6 +5,7 @@ from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 
 from .const import (
     ATTR_MANUFACTURER,
+    CONF_CONTROLLER,
     DOMAIN as UNIFI_DOMAIN,
     LOGGER,
     UNIFI_WIRELESS_CLIENTS,
@@ -28,10 +29,14 @@ async def async_setup_entry(hass, config_entry):
     """Set up the UniFi component."""
     hass.data.setdefault(UNIFI_DOMAIN, {})
 
+    # Flat configuration was introduced with 2021.3
+    await async_flatten_entry_data(hass, config_entry)
+
     controller = UniFiController(hass, config_entry)
     if not await controller.async_setup():
         return False
 
+    # Unique ID was introduced with 2021.3
     if config_entry.unique_id is None:
         hass.config_entries.async_update_entry(
             config_entry, unique_id=controller.site_id
@@ -64,6 +69,17 @@ async def async_unload_entry(hass, config_entry):
     return await controller.async_reset()
 
 
+async def async_flatten_entry_data(hass, config_entry):
+    """Simpler configuration structure for entry data.
+
+    Keep controller key layer in case user rollbacks.
+    """
+
+    data: dict = {**config_entry.data, **config_entry.data[CONF_CONTROLLER]}
+    if config_entry.data != data:
+        hass.config_entries.async_update_entry(config_entry, data=data)
+
+
 class UnifiWirelessClients:
     """Class to store clients known to be wireless.
 

homeassistant/components/unifi/config_flow.py

@@ -73,7 +73,6 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
         self.site_ids = {}
         self.site_names = {}
         self.reauth_config_entry = None
-        self.reauth_config = {}
         self.reauth_schema = {}
 
     async def async_step_user(self, user_input=None):
@@ -92,7 +91,16 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
             }
 
             try:
-                controller = await get_controller(self.hass, **self.config)
+                controller = await get_controller(
+                    self.hass,
+                    host=self.config[CONF_HOST],
+                    username=self.config[CONF_USERNAME],
+                    password=self.config[CONF_PASSWORD],
+                    port=self.config[CONF_PORT],
+                    site=self.config[CONF_SITE_ID],
+                    verify_ssl=self.config[CONF_VERIFY_SSL],
+                )
+
                 sites = await controller.sites()
 
             except AuthenticationRequired:
@@ -143,7 +151,8 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
 
             unique_id = user_input[CONF_SITE_ID]
             self.config[CONF_SITE_ID] = self.site_ids[unique_id]
-            data = {CONF_CONTROLLER: self.config}
+            # Backwards compatible config
+            self.config[CONF_CONTROLLER] = self.config.copy()
 
             config_entry = await self.async_set_unique_id(unique_id)
             abort_reason = "configuration_updated"
@@ -160,12 +169,14 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
                 if controller and controller.available:
                     return self.async_abort(reason="already_configured")
 
-                self.hass.config_entries.async_update_entry(config_entry, data=data)
+                self.hass.config_entries.async_update_entry(
+                    config_entry, data=self.config
+                )
                 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]
-            return self.async_create_entry(title=site_nice_name, data=data)
+            return self.async_create_entry(title=site_nice_name, data=self.config)
 
         if len(self.site_names) == 1:
             return await self.async_step_site(
@@ -183,21 +194,20 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
     async def async_step_reauth(self, config_entry: dict):
         """Trigger a reauthentication flow."""
         self.reauth_config_entry = config_entry
-        self.reauth_config = config_entry.data[CONF_CONTROLLER]
 
         # pylint: disable=no-member # https://github.com/PyCQA/pylint/issues/3167
         self.context["title_placeholders"] = {
-            CONF_HOST: self.reauth_config[CONF_HOST],
+            CONF_HOST: config_entry.data[CONF_HOST],
             CONF_SITE_ID: config_entry.title,
         }
 
         self.reauth_schema = {
-            vol.Required(CONF_HOST, default=self.reauth_config[CONF_HOST]): str,
+            vol.Required(CONF_HOST, default=config_entry.data[CONF_HOST]): str,
-            vol.Required(CONF_USERNAME, default=self.reauth_config[CONF_USERNAME]): str,
+            vol.Required(CONF_USERNAME, default=config_entry.data[CONF_USERNAME]): str,
             vol.Required(CONF_PASSWORD): str,
-            vol.Required(CONF_PORT, default=self.reauth_config[CONF_PORT]): int,
+            vol.Required(CONF_PORT, default=config_entry.data[CONF_PORT]): int,
             vol.Required(
-                CONF_VERIFY_SSL, default=self.reauth_config[CONF_VERIFY_SSL]
+                CONF_VERIFY_SSL, default=config_entry.data[CONF_VERIFY_SSL]
             ): bool,
         }
 
@@ -217,7 +227,7 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
             return self.async_abort(reason="already_configured")
 
         await self.async_set_unique_id(mac_address)
-        self._abort_if_unique_id_configured(updates={CONF_HOST: self.config[CONF_HOST]})
+        self._abort_if_unique_id_configured(updates=self.config)
 
         # pylint: disable=no-member
         self.context["title_placeholders"] = {
@@ -234,9 +244,7 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
     def _host_already_configured(self, host):
         """See if we already have a UniFi entry matching the host."""
         for entry in self._async_current_entries():
-            if not entry.data or CONF_CONTROLLER not in entry.data:
+            if entry.data.get(CONF_HOST) == host:
-                continue
-            if entry.data[CONF_CONTROLLER][CONF_HOST] == host:
                 return True
         return False
 

homeassistant/components/unifi/controller.py

@@ -29,7 +29,13 @@ from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
 from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.config_entries import SOURCE_REAUTH
-from homeassistant.const import CONF_HOST
+from homeassistant.const import (
+    CONF_HOST,
+    CONF_PASSWORD,
+    CONF_PORT,
+    CONF_USERNAME,
+    CONF_VERIFY_SSL,
+)
 from homeassistant.core import callback
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import aiohttp_client
@@ -41,7 +47,6 @@ from .const import (
     CONF_ALLOW_BANDWIDTH_SENSORS,
     CONF_ALLOW_UPTIME_SENSORS,
     CONF_BLOCK_CLIENT,
-    CONF_CONTROLLER,
     CONF_DETECTION_TIME,
     CONF_DPI_RESTRICTIONS,
     CONF_IGNORE_WIRED_BUG,
@@ -161,12 +166,12 @@ class UniFiController:
     @property
     def host(self):
         """Return the host of this controller."""
-        return self.config_entry.data[CONF_CONTROLLER][CONF_HOST]
+        return self.config_entry.data[CONF_HOST]
 
     @property
     def site(self):
         """Return the site of this config entry."""
-        return self.config_entry.data[CONF_CONTROLLER][CONF_SITE_ID]
+        return self.config_entry.data[CONF_SITE_ID]
 
     @property
     def site_name(self):
@@ -299,7 +304,12 @@ class UniFiController:
         try:
             self.api = await get_controller(
                 self.hass,
-                **self.config_entry.data[CONF_CONTROLLER],
+                host=self.config_entry.data[CONF_HOST],
+                username=self.config_entry.data[CONF_USERNAME],
+                password=self.config_entry.data[CONF_PASSWORD],
+                port=self.config_entry.data[CONF_PORT],
+                site=self.config_entry.data[CONF_SITE_ID],
+                verify_ssl=self.config_entry.data[CONF_VERIFY_SSL],
                 async_callback=self.async_unifi_signalling_callback,
             )
             await self.api.initialize()

tests/components/unifi/test_config_flow.py

@@ -134,6 +134,12 @@ async def test_flow_works(hass, aioclient_mock, mock_discovery):
     assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
     assert result["title"] == "Site name"
     assert result["data"] == {
+        CONF_HOST: "1.2.3.4",
+        CONF_USERNAME: "username",
+        CONF_PASSWORD: "password",
+        CONF_PORT: 1234,
+        CONF_SITE_ID: "site_id",
+        CONF_VERIFY_SSL: True,
         CONF_CONTROLLER: {
             CONF_HOST: "1.2.3.4",
             CONF_USERNAME: "username",
@@ -141,7 +147,7 @@ async def test_flow_works(hass, aioclient_mock, mock_discovery):
             CONF_PORT: 1234,
             CONF_SITE_ID: "site_id",
             CONF_VERIFY_SSL: True,
-        }
+        },
     }
 
 
@@ -241,16 +247,12 @@ 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,
+        domain=UNIFI_DOMAIN, data={"host": "1.2.3.4", "site": "office"}, unique_id="2"
-        data={"controller": {"host": "1.2.3.4", "site": "office"}},
-        unique_id="2",
     )
     entry.add_to_hass(hass)
 
     entry = MockConfigEntry(
-        domain=UNIFI_DOMAIN,
+        domain=UNIFI_DOMAIN, data={"host": "1.2.3.4", "site": "site_id"}, unique_id="1"
-        data={"controller": {"host": "1.2.3.4", "site": "site_id"}},
-        unique_id="1",
     )
     entry.add_to_hass(hass)
 
@@ -399,9 +401,9 @@ async def test_reauth_flow_update_configuration(hass, aioclient_mock):
 
     assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT
     assert result["reason"] == "reauth_successful"
-    assert config_entry.data[CONF_CONTROLLER][CONF_HOST] == "1.2.3.4"
+    assert config_entry.data[CONF_HOST] == "1.2.3.4"
-    assert config_entry.data[CONF_CONTROLLER][CONF_USERNAME] == "new_name"
+    assert config_entry.data[CONF_USERNAME] == "new_name"
-    assert config_entry.data[CONF_CONTROLLER][CONF_PASSWORD] == "new_pass"
+    assert config_entry.data[CONF_PASSWORD] == "new_pass"
 
 
 async def test_advanced_option_flow(hass, aioclient_mock):
@@ -544,7 +546,7 @@ async def test_form_ssdp_aborts_if_host_already_exists(hass):
     await setup.async_setup_component(hass, "persistent_notification", {})
     entry = MockConfigEntry(
         domain=UNIFI_DOMAIN,
-        data={"controller": {"host": "192.168.208.1", "site": "site_id"}},
+        data={"host": "192.168.208.1", "site": "site_id"},
     )
     entry.add_to_hass(hass)
     result = await hass.config_entries.flow.async_init(

tests/components/unifi/test_controller.py

@@ -66,7 +66,7 @@ CONTROLLER_DATA = {
     CONF_VERIFY_SSL: False,
 }
 
-ENTRY_CONFIG = {CONF_CONTROLLER: CONTROLLER_DATA}
+ENTRY_CONFIG = {**CONTROLLER_DATA, CONF_CONTROLLER: CONTROLLER_DATA}
 ENTRY_OPTIONS = {}
 
 CONFIGURATION = []
@@ -167,6 +167,7 @@ async def setup_unifi_integration(
         options=deepcopy(options),
         entry_id=1,
         unique_id="1",
+        version=1,
     )
     config_entry.add_to_hass(hass)
 
@@ -178,8 +179,8 @@ async def setup_unifi_integration(
     if aioclient_mock:
         mock_default_unifi_requests(
             aioclient_mock,
-            host=config_entry.data[CONF_CONTROLLER][CONF_HOST],
+            host=config_entry.data[CONF_HOST],
-            site_id=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID],
+            site_id=config_entry.data[CONF_SITE_ID],
             sites=sites,
             description=site_description,
             clients_response=clients_response,

tests/components/unifi/test_init.py

@@ -2,10 +2,11 @@
 from unittest.mock import AsyncMock, Mock, patch
 
 from homeassistant.components import unifi
-from homeassistant.components.unifi.const import DOMAIN as UNIFI_DOMAIN
+from homeassistant.components.unifi import async_flatten_entry_data
+from homeassistant.components.unifi.const import CONF_CONTROLLER, DOMAIN as UNIFI_DOMAIN
 from homeassistant.setup import async_setup_component
 
-from .test_controller import setup_unifi_integration
+from .test_controller import CONTROLLER_DATA, ENTRY_CONFIG, setup_unifi_integration
 
 from tests.common import MockConfigEntry, mock_coro
 
@@ -35,17 +36,9 @@ async def test_controller_no_mac(hass):
     """Test that configured options for a host are loaded via config entry."""
     entry = MockConfigEntry(
         domain=UNIFI_DOMAIN,
-        data={
+        data=ENTRY_CONFIG,
-            "controller": {
-                "host": "0.0.0.0",
-                "username": "user",
-                "password": "pass",
-                "port": 80,
-                "site": "default",
-                "verify_ssl": True,
-            },
-        },
         unique_id="1",
+        version=1,
     )
     entry.add_to_hass(hass)
     mock_registry = Mock()
@@ -64,6 +57,17 @@ async def test_controller_no_mac(hass):
     assert len(mock_registry.mock_calls) == 0
 
 
+async def test_flatten_entry_data(hass):
+    """Verify entry data can be flattened."""
+    entry = MockConfigEntry(
+        domain=UNIFI_DOMAIN,
+        data={CONF_CONTROLLER: CONTROLLER_DATA},
+    )
+    await async_flatten_entry_data(hass, entry)
+
+    assert entry.data == ENTRY_CONFIG
+
+
 async def test_unload_entry(hass, aioclient_mock):
     """Test being able to unload an entry."""
     config_entry = await setup_unifi_integration(hass, aioclient_mock)