evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Introduce reauthentication flow to UniFi integration (#45360)

Browse source 
* Improve site selection

* Reauth flow and tests
Add **kwargs to mock_aiohttp_client create_session to support inputting verify_ssl and cookie_jar

* Update homeassistant/components/unifi/config_flow.py

Co-authored-by: J. Nick Koston <nick@koston.org>

* Minor improvements

* Improve coverage

Co-authored-by: J. Nick Koston <nick@koston.org>
This commit is contained in:
Robert Svensson 2021-01-20 22:10:40 +01:00 • committed by GitHub
parent 7ff02fe8d4
commit da4404e8cf
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 157 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

82
homeassistant/components/unifi/config_flow.py
View file

@ -66,8 +66,10 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
def __init__(self):
"""Initialize the UniFi flow."""
self.config = None
self.desc = None
self.sites = None
self.reauth_config_entry = {}
self.reauth_config = {}
self.reauth_schema = {}
async def async_step_user(self, user_input=None):
"""Handle a flow initialized by the user."""
@ -87,7 +89,13 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
controller = await get_controller(self.hass, **self.config)
self.sites = await controller.sites()
sites = await controller.sites()
self.sites = {site["name"]: site["desc"] for site in sites.values()}
if self.reauth_config.get(CONF_SITE_ID) in self.sites:
return await self.async_step_site(
{CONF_SITE_ID: self.reauth_config[CONF_SITE_ID]}
)
return await self.async_step_site()
@ -108,17 +116,17 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
if await async_discover_unifi(self.hass):
host = "unifi"
data = self.reauth_schema or {
vol.Required(CONF_HOST, default=host): str,
vol.Required(CONF_USERNAME): str,
vol.Required(CONF_PASSWORD): str,
vol.Optional(CONF_PORT, default=DEFAULT_PORT): int,
vol.Optional(CONF_VERIFY_SSL, default=DEFAULT_VERIFY_SSL): bool,
}
return self.async_show_form(
step_id="user",
data_schema=vol.Schema(
{
vol.Required(CONF_HOST, default=host): str,
vol.Required(CONF_USERNAME): str,
vol.Required(CONF_PASSWORD): str,
vol.Optional(CONF_PORT, default=DEFAULT_PORT): int,
vol.Optional(CONF_VERIFY_SSL, default=DEFAULT_VERIFY_SSL): bool,
}
),
data_schema=vol.Schema(data),
errors=errors,
)
@ -128,12 +136,17 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
if user_input is not None:
try:
desc = user_input.get(CONF_SITE_ID, self.desc)
self.config[CONF_SITE_ID] = user_input[CONF_SITE_ID]
data = {CONF_CONTROLLER: self.config}
for site in self.sites.values():
if desc == site["desc"]:
self.config[CONF_SITE_ID] = site["name"]
break
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 entry in self._async_current_entries():
controller = entry.data[CONF_CONTROLLER]
@ -143,27 +156,44 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
):
raise AlreadyConfigured
data = {CONF_CONTROLLER: self.config}
return self.async_create_entry(title=desc, data=data)
site_nice_name = self.sites[self.config[CONF_SITE_ID]]
return self.async_create_entry(title=site_nice_name, data=data)
except AlreadyConfigured:
return self.async_abort(reason="already_configured")
if len(self.sites) == 1:
self.desc = next(iter(self.sites.values()))["desc"]
return await self.async_step_site(user_input={})
sites = []
for site in self.sites.values():
sites.append(site["desc"])
return await self.async_step_site({CONF_SITE_ID: next(iter(self.sites))})
return self.async_show_form(
step_id="site",
data_schema=vol.Schema({vol.Required(CONF_SITE_ID): vol.In(sites)}),
data_schema=vol.Schema({vol.Required(CONF_SITE_ID): vol.In(self.sites)}),
errors=errors,
)
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_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_PASSWORD): str,
vol.Required(CONF_PORT, default=self.reauth_config[CONF_PORT]): int,
vol.Required(
CONF_VERIFY_SSL, default=self.reauth_config[CONF_VERIFY_SSL]
): bool,
}
return await self.async_step_user()
class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
"""Handle Unifi options."""

31
homeassistant/components/unifi/controller.py
View file

@ -28,6 +28,7 @@ import async_timeout
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.core import callback
from homeassistant.exceptions import ConfigEntryNotReady
@ -343,8 +344,14 @@ class UniFiController:
except CannotConnect as err:
raise ConfigEntryNotReady from err
except Exception as err: # pylint: disable=broad-except
LOGGER.error("Unknown error connecting with UniFi controller: %s", err)
except AuthenticationRequired:
self.hass.async_create_task(
self.hass.config_entries.flow.async_init(
UNIFI_DOMAIN,
context={"source": SOURCE_REAUTH},
data=self.config_entry,
)
)
return False
# Restore clients that is not a part of active clients list.
@ -419,7 +426,13 @@ class UniFiController:
@staticmethod
async def async_config_entry_updated(hass, config_entry) -> None:
"""Handle signals of config entry being updated."""
"""Handle signals of config entry being updated.
If config entry is updated due to reauth flow
the entry might already have been reset and thus is not available.
"""
if config_entry.entry_id not in hass.data[UNIFI_DOMAIN]:
return
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
async_dispatcher_send(hass, controller.signal_options_update)
@ -510,7 +523,7 @@ async def get_controller(
return controller
except aiounifi.Unauthorized as err:
LOGGER.warning("Connected to UniFi at %s but not registered.", host)
LOGGER.warning("Connected to UniFi at %s but not registered: %s", host, err)
raise AuthenticationRequired from err
except (
@ -519,9 +532,13 @@ async def get_controller(
aiounifi.ServiceUnavailable,
aiounifi.RequestError,
) as err:
LOGGER.error("Error connecting to the UniFi controller at %s", host)
LOGGER.error("Error connecting to the UniFi controller at %s: %s", host, err)
raise CannotConnect from err
except aiounifi.AiounifiException as err:
LOGGER.exception("Unknown UniFi communication error occurred")
except aiounifi.LoginRequired as err:
LOGGER.warning("Connected to UniFi at %s but login required: %s", host, err)
raise AuthenticationRequired from err
except aiounifi.AiounifiException as err:
LOGGER.exception("Unknown UniFi communication error occurred: %s", err)
raise AuthenticationRequired from err

4
homeassistant/components/unifi/strings.json
View file

@ -1,5 +1,6 @@
{
"config": {
"flow_title": "{site} ({host})",
"step": {
"user": {
"title": "Set up UniFi Controller",
@ -19,7 +20,8 @@
"unknown_client_mac": "No client available on that MAC address"
},
"abort": {
"already_configured": "Controller site is already configured"
"already_configured": "Controller site is already configured",
"reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]"
}
},
"options": {

4
homeassistant/components/unifi/translations/en.json
View file

@ -1,13 +1,15 @@
{
"config": {
"abort": {
"already_configured": "Controller site is already configured"
"already_configured": "Controller site is already configured",
"reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]"
},
"error": {
"faulty_credentials": "Invalid authentication",
"service_unavailable": "Failed to connect",
"unknown_client_mac": "No client available on that MAC address"
},
"flow_title": "{site} ({host})",
"step": {
"user": {
"data": {

53
tests/components/unifi/test_config_flow.py
View file

@ -20,6 +20,7 @@ from homeassistant.components.unifi.const import (
CONF_TRACK_WIRED_CLIENTS,
DOMAIN as UNIFI_DOMAIN,
)
from homeassistant.config_entries import SOURCE_REAUTH, SOURCE_USER
from homeassistant.const import (
CONF_HOST,
CONF_PASSWORD,
@ -184,8 +185,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": "site name"})
assert result["data_schema"]({"site": "site2 name"})
assert result["data_schema"]({"site": "default"})
assert result["data_schema"]({"site": "site2"})
async def test_flow_fails_site_already_configured(hass, aioclient_mock):
@ -314,6 +315,54 @@ async def test_flow_fails_unknown_problem(hass, aioclient_mock):
assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT
async def test_reauth_flow_update_configuration(hass, aioclient_mock):
"""Verify reauth flow can update controller configuration."""
controller = await setup_unifi_integration(hass)
result = await hass.config_entries.flow.async_init(
UNIFI_DOMAIN,
context={"source": SOURCE_REAUTH},
data=controller.config_entry,
)
assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
assert result["step_id"] == SOURCE_USER
aioclient_mock.get("https://1.2.3.4:1234", status=302)
aioclient_mock.post(
"https://1.2.3.4:1234/api/login",
json={"data": "login successful", "meta": {"rc": "ok"}},
headers={"content-type": CONTENT_TYPE_JSON},
)
aioclient_mock.get(
"https://1.2.3.4:1234/api/self/sites",
json={
"data": [{"desc": "Site name", "name": "site_id", "role": "admin"}],
"meta": {"rc": "ok"},
},
headers={"content-type": CONTENT_TYPE_JSON},
)
result = await hass.config_entries.flow.async_configure(
result["flow_id"],
user_input={
CONF_HOST: "1.2.3.4",
CONF_USERNAME: "new_name",
CONF_PASSWORD: "new_pass",
CONF_PORT: 1234,
CONF_VERIFY_SSL: True,
},
)
assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT
assert result["reason"] == "reauth_successful"
assert controller.host == "1.2.3.4"
assert controller.config_entry.data[CONF_CONTROLLER][CONF_USERNAME] == "new_name"
assert controller.config_entry.data[CONF_CONTROLLER][CONF_PASSWORD] == "new_pass"
async def test_advanced_option_flow(hass):
"""Test advanced config flow options."""
controller = await setup_unifi_integration(

19
tests/components/unifi/test_controller.py
View file

@ -222,6 +222,17 @@ async def test_controller_not_accessible(hass):
assert hass.data[UNIFI_DOMAIN] == {}
async def test_controller_trigger_reauth_flow(hass):
"""Failed authentication trigger a reauthentication flow."""
with patch(
"homeassistant.components.unifi.controller.get_controller",
side_effect=AuthenticationRequired,
), patch.object(hass.config_entries.flow, "async_init") as mock_flow_init:
await setup_unifi_integration(hass)
mock_flow_init.assert_called_once()
assert hass.data[UNIFI_DOMAIN] == {}
async def test_controller_unknown_error(hass):
"""Unknown errors are handled."""
with patch(
@ -319,6 +330,14 @@ async def test_get_controller_controller_unavailable(hass):
await get_controller(hass, **CONTROLLER_DATA)
async def test_get_controller_login_required(hass):
"""Check that get_controller can handle unknown errors."""
with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(
"aiounifi.Controller.login", side_effect=aiounifi.LoginRequired
), pytest.raises(AuthenticationRequired):
await get_controller(hass, **CONTROLLER_DATA)
async def test_get_controller_unknown_error(hass):
"""Check that get_controller can handle unknown errors."""
with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(

2
tests/test_util/aiohttp.py
View file

@ -276,7 +276,7 @@ def mock_aiohttp_client():
"""Context manager to mock aiohttp client."""
mocker = AiohttpClientMocker()
def create_session(hass, *args):
def create_session(hass, *args, **kwargs):
session = mocker.create_session(hass.loop)
async def close_session(event):