diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index d02f3f49a5e..50e578b1dae 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/homeassistant/components/unifi/config_flow.py @@ -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") diff --git a/tests/components/unifi/conftest.py b/tests/components/unifi/conftest.py index e21c458386f..e2b77ac1ed1 100644 --- a/tests/components/unifi/conftest.py +++ b/tests/components/unifi/conftest.py @@ -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 diff --git a/tests/components/unifi/test_config_flow.py b/tests/components/unifi/test_config_flow.py index 321cbdfd9e8..e774c5a551d 100644 --- a/tests/components/unifi/test_config_flow.py +++ b/tests/components/unifi/test_config_flow.py @@ -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