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
This commit is contained in:
Robert Svensson 2021-02-06 21:32:18 +01:00 committed by GitHub
parent cefde8721d
commit 618fcda821
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 87 additions and 46 deletions

View file

@ -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.

View file

@ -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

View file

@ -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()

View file

@ -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(

View file

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

View file

@ -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)