diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index 5f19a01ce45..8d24a9b642f 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/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. diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index 1d89215dc89..8e83f53d198 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/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_USERNAME, default=self.reauth_config[CONF_USERNAME]): str, + vol.Required(CONF_HOST, default=config_entry.data[CONF_HOST]): 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: - continue - if entry.data[CONF_CONTROLLER][CONF_HOST] == host: + if entry.data.get(CONF_HOST) == host: return True return False diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 5d5e679e75e..128f0107984 100644 --- a/homeassistant/components/unifi/controller.py +++ b/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() diff --git a/tests/components/unifi/test_config_flow.py b/tests/components/unifi/test_config_flow.py index 096e6ba7791..a28f5f5f7c5 100644 --- a/tests/components/unifi/test_config_flow.py +++ b/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, - data={"controller": {"host": "1.2.3.4", "site": "office"}}, - unique_id="2", + domain=UNIFI_DOMAIN, data={"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"}}, - unique_id="1", + domain=UNIFI_DOMAIN, data={"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_CONTROLLER][CONF_USERNAME] == "new_name" - assert config_entry.data[CONF_CONTROLLER][CONF_PASSWORD] == "new_pass" + assert config_entry.data[CONF_HOST] == "1.2.3.4" + assert config_entry.data[CONF_USERNAME] == "new_name" + 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( diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index 3ecd44b3db7..00865b4e910 100644 --- a/tests/components/unifi/test_controller.py +++ b/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], - site_id=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID], + host=config_entry.data[CONF_HOST], + site_id=config_entry.data[CONF_SITE_ID], sites=sites, description=site_description, clients_response=clients_response, diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py index 9de8b0a0990..6d8b894fc34 100644 --- a/tests/components/unifi/test_init.py +++ b/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={ - "controller": { - "host": "0.0.0.0", - "username": "user", - "password": "pass", - "port": 80, - "site": "default", - "verify_ssl": True, - }, - }, + data=ENTRY_CONFIG, 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)