Adjust config-flow type hints in unifi (#72411)

* Adjust config-flow type hints in unifi

* Use mapping

* Use mapping

* Fix tests

* Fix tests

* Revert "Use mapping"

This reverts commit 126fedc848.
This commit is contained in:
epenet 2022-05-24 21:42:11 +02:00 committed by GitHub
parent 2e36a79357
commit f33151ff8b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 39 additions and 19 deletions

View file

@ -5,7 +5,11 @@ Discovery of UniFi Network instances hosted on UDM and UDM Pro devices
through SSDP. Reauthentication when issue with credentials are reported.
Configuration of options through options flow.
"""
from __future__ import annotations
from collections.abc import Mapping
import socket
from typing import Any
from urllib.parse import urlparse
import voluptuous as vol
@ -19,7 +23,7 @@ from homeassistant.const import (
CONF_USERNAME,
CONF_VERIFY_SSL,
)
from homeassistant.core import callback
from homeassistant.core import HomeAssistant, callback
from homeassistant.data_entry_flow import FlowResult
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.device_registry import format_mac
@ -63,11 +67,13 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
@staticmethod
@callback
def async_get_options_flow(config_entry):
def async_get_options_flow(
config_entry: config_entries.ConfigEntry,
) -> UnifiOptionsFlowHandler:
"""Get the options flow for this handler."""
return UnifiOptionsFlowHandler(config_entry)
def __init__(self):
def __init__(self) -> None:
"""Initialize the UniFi Network flow."""
self.config = {}
self.site_ids = {}
@ -75,7 +81,9 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
self.reauth_config_entry = None
self.reauth_schema = {}
async def async_step_user(self, user_input=None):
async def async_step_user(
self, user_input: dict[str, Any] | None = None
) -> FlowResult:
"""Handle a flow initialized by the user."""
errors = {}
@ -123,7 +131,7 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
return await self.async_step_site()
if not (host := self.config.get(CONF_HOST, "")) and await async_discover_unifi(
if not (host := self.config.get(CONF_HOST, "")) and await _async_discover_unifi(
self.hass
):
host = "unifi"
@ -144,7 +152,9 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
errors=errors,
)
async def async_step_site(self, user_input=None):
async def async_step_site(
self, user_input: dict[str, Any] | None = None
) -> FlowResult:
"""Select site to control."""
errors = {}
@ -192,7 +202,7 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
errors=errors,
)
async def async_step_reauth(self, data: dict):
async def async_step_reauth(self, data: Mapping[str, Any]) -> FlowResult:
"""Trigger a reauthentication flow."""
config_entry = self.hass.config_entries.async_get_entry(
self.context["entry_id"]
@ -248,13 +258,15 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
"""Handle Unifi Network options."""
def __init__(self, config_entry):
def __init__(self, config_entry: config_entries.ConfigEntry) -> None:
"""Initialize UniFi Network options flow."""
self.config_entry = config_entry
self.options = dict(config_entry.options)
self.controller = None
async def async_step_init(self, user_input=None):
async def async_step_init(
self, user_input: dict[str, Any] | None = None
) -> FlowResult:
"""Manage the UniFi Network options."""
if self.config_entry.entry_id not in self.hass.data[UNIFI_DOMAIN]:
return self.async_abort(reason="integration_not_setup")
@ -266,7 +278,9 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
return await self.async_step_simple_options()
async def async_step_simple_options(self, user_input=None):
async def async_step_simple_options(
self, user_input: dict[str, Any] | None = None
) -> FlowResult:
"""For users without advanced settings enabled."""
if user_input is not None:
self.options.update(user_input)
@ -299,7 +313,9 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
last_step=True,
)
async def async_step_device_tracker(self, user_input=None):
async def async_step_device_tracker(
self, user_input: dict[str, Any] | None = None
) -> FlowResult:
"""Manage the device tracker options."""
if user_input is not None:
self.options.update(user_input)
@ -359,7 +375,9 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
last_step=False,
)
async def async_step_client_control(self, user_input=None):
async def async_step_client_control(
self, user_input: dict[str, Any] | None = None
) -> FlowResult:
"""Manage configuration of network access controlled clients."""
errors = {}
@ -403,7 +421,9 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
last_step=False,
)
async def async_step_statistics_sensors(self, user_input=None):
async def async_step_statistics_sensors(
self, user_input: dict[str, Any] | None = None
) -> FlowResult:
"""Manage the statistics sensors options."""
if user_input is not None:
self.options.update(user_input)
@ -426,12 +446,12 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
last_step=True,
)
async def _update_options(self):
async def _update_options(self) -> FlowResult:
"""Update config entry options."""
return self.async_create_entry(title="", data=self.options)
async def async_discover_unifi(hass):
async def _async_discover_unifi(hass: HomeAssistant) -> str | None:
"""Discover UniFi Network address."""
try:
return await hass.async_add_executor_job(socket.gethostbyname, "unifi")

View file

@ -34,7 +34,7 @@ def mock_unifi_websocket():
def mock_discovery():
"""No real network traffic allowed."""
with patch(
"homeassistant.components.unifi.config_flow.async_discover_unifi",
"homeassistant.components.unifi.config_flow._async_discover_unifi",
return_value=None,
) as mock:
yield mock

View file

@ -7,7 +7,7 @@ import aiounifi
from homeassistant import config_entries, data_entry_flow
from homeassistant.components import ssdp
from homeassistant.components.unifi.config_flow import async_discover_unifi
from homeassistant.components.unifi.config_flow import _async_discover_unifi
from homeassistant.components.unifi.const import (
CONF_ALLOW_BANDWIDTH_SENSORS,
CONF_ALLOW_UPTIME_SENSORS,
@ -686,10 +686,10 @@ async def test_form_ssdp_gets_form_with_ignored_entry(hass):
async def test_discover_unifi_positive(hass):
"""Verify positive run of UniFi discovery."""
with patch("socket.gethostbyname", return_value=True):
assert await async_discover_unifi(hass)
assert await _async_discover_unifi(hass)
async def test_discover_unifi_negative(hass):
"""Verify negative run of UniFi discovery."""
with patch("socket.gethostbyname", side_effect=socket.gaierror):
assert await async_discover_unifi(hass) is None
assert await _async_discover_unifi(hass) is None