From b68c90d59a93953baf9dc43c8a3c55bc1d0488f3 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Fri, 6 Sep 2024 15:16:47 +0200 Subject: [PATCH] Improve config flow type hints in vulcan (#125308) * Improve config flow type hints in vulcan * Adjust tests --- .../components/vulcan/config_flow.py | 56 +++++++++++++------ homeassistant/components/vulcan/register.py | 4 +- tests/components/vulcan/test_config_flow.py | 4 +- 3 files changed, 43 insertions(+), 21 deletions(-) diff --git a/homeassistant/components/vulcan/config_flow.py b/homeassistant/components/vulcan/config_flow.py index 5938e4ce690..f02adba9f75 100644 --- a/homeassistant/components/vulcan/config_flow.py +++ b/homeassistant/components/vulcan/config_flow.py @@ -2,7 +2,7 @@ from collections.abc import Mapping import logging -from typing import Any +from typing import TYPE_CHECKING, Any from aiohttp import ClientConnectionError import voluptuous as vol @@ -16,6 +16,7 @@ from vulcan import ( UnauthorizedCertificateException, Vulcan, ) +from vulcan.model import Student from homeassistant.config_entries import ConfigFlow, ConfigFlowResult from homeassistant.const import CONF_PIN, CONF_REGION, CONF_TOKEN @@ -38,11 +39,12 @@ class VulcanFlowHandler(ConfigFlow, domain=DOMAIN): VERSION = 1 + account: Account + keystore: Keystore + def __init__(self) -> None: """Initialize config flow.""" - self.account = None - self.keystore = None - self.students = None + self.students: list[Student] | None = None async def async_step_user( self, user_input: dict[str, Any] | None = None @@ -53,13 +55,16 @@ class VulcanFlowHandler(ConfigFlow, domain=DOMAIN): return await self.async_step_auth() - async def async_step_auth(self, user_input=None, errors=None): + async def async_step_auth( + self, + user_input: dict[str, str] | None = None, + errors: dict[str, str] | None = None, + ) -> ConfigFlowResult: """Authorize integration.""" if user_input is not None: try: credentials = await register( - self.hass, user_input[CONF_TOKEN], user_input[CONF_REGION], user_input[CONF_PIN], @@ -107,16 +112,20 @@ class VulcanFlowHandler(ConfigFlow, domain=DOMAIN): errors=errors, ) - async def async_step_select_student(self, user_input=None): + async def async_step_select_student( + self, user_input: dict[str, str] | None = None + ) -> ConfigFlowResult: """Allow user to select student.""" - errors = {} - students = {} + errors: dict[str, str] = {} + students: dict[str, str] = {} if self.students is not None: for student in self.students: students[str(student.pupil.id)] = ( f"{student.pupil.first_name} {student.pupil.last_name}" ) if user_input is not None: + if TYPE_CHECKING: + assert self.keystore is not None student_id = user_input["student"] await self.async_set_unique_id(str(student_id)) self._abort_if_unique_id_configured() @@ -135,17 +144,25 @@ class VulcanFlowHandler(ConfigFlow, domain=DOMAIN): errors=errors, ) - async def async_step_select_saved_credentials(self, user_input=None, errors=None): + async def async_step_select_saved_credentials( + self, + user_input: dict[str, str] | None = None, + errors: dict[str, str] | None = None, + ) -> ConfigFlowResult: """Allow user to select saved credentials.""" - credentials = {} + credentials: dict[str, Any] = {} for entry in self.hass.config_entries.async_entries(DOMAIN): credentials[entry.entry_id] = entry.data["account"]["UserName"] if user_input is not None: - entry = self.hass.config_entries.async_get_entry(user_input["credentials"]) - keystore = Keystore.load(entry.data["keystore"]) - account = Account.load(entry.data["account"]) + existing_entry = self.hass.config_entries.async_get_entry( + user_input["credentials"] + ) + if TYPE_CHECKING: + assert existing_entry is not None + keystore = Keystore.load(existing_entry.data["keystore"]) + account = Account.load(existing_entry.data["account"]) client = Vulcan(keystore, account, async_get_clientsession(self.hass)) try: students = await client.get_students() @@ -189,12 +206,14 @@ class VulcanFlowHandler(ConfigFlow, domain=DOMAIN): errors=errors, ) - async def async_step_add_next_config_entry(self, user_input=None): + async def async_step_add_next_config_entry( + self, user_input: dict[str, bool] | None = None + ) -> ConfigFlowResult: """Flow initialized when user is adding next entry of that integration.""" existing_entries = self.hass.config_entries.async_entries(DOMAIN) - errors = {} + errors: dict[str, str] = {} if user_input is not None: if not user_input["use_saved_credentials"]: @@ -248,13 +267,14 @@ class VulcanFlowHandler(ConfigFlow, domain=DOMAIN): """Perform reauth upon an API authentication error.""" return await self.async_step_reauth_confirm() - async def async_step_reauth_confirm(self, user_input=None): + async def async_step_reauth_confirm( + self, user_input: dict[str, str] | None = None + ) -> ConfigFlowResult: """Reauthorize integration.""" errors = {} if user_input is not None: try: credentials = await register( - self.hass, user_input[CONF_TOKEN], user_input[CONF_REGION], user_input[CONF_PIN], diff --git a/homeassistant/components/vulcan/register.py b/homeassistant/components/vulcan/register.py index 67cceb8d7b8..a3dec97f622 100644 --- a/homeassistant/components/vulcan/register.py +++ b/homeassistant/components/vulcan/register.py @@ -1,9 +1,11 @@ """Support for register Vulcan account.""" +from typing import Any + from vulcan import Account, Keystore -async def register(hass, token, symbol, pin): +async def register(token: str, symbol: str, pin: str) -> dict[str, Any]: """Register integration and save credentials.""" keystore = await Keystore.create(device_model="Home Assistant") account = await Account.register(keystore, token, symbol, pin) diff --git a/tests/components/vulcan/test_config_flow.py b/tests/components/vulcan/test_config_flow.py index a72e77b32e8..a51d9727126 100644 --- a/tests/components/vulcan/test_config_flow.py +++ b/tests/components/vulcan/test_config_flow.py @@ -310,7 +310,7 @@ async def test_multiple_config_entries( unique_id="123456", data=json.loads(load_fixture("fake_config_entry_data.json", "vulcan")), ).add_to_hass(hass) - await register.register(hass, "token", "region", "000000") + await register.register("token", "region", "000000") result = await hass.config_entries.flow.async_init( const.DOMAIN, context={"source": config_entries.SOURCE_USER} ) @@ -703,7 +703,7 @@ async def test_student_already_exists( | {"student_id": "0"}, ).add_to_hass(hass) - await register.register(hass, "token", "region", "000000") + await register.register("token", "region", "000000") result = await hass.config_entries.flow.async_init( const.DOMAIN, context={"source": config_entries.SOURCE_USER}