Use new reauth helpers in unifi (#128837)

* Use new reauth helpers in unifi * Apply suggestions from code review Co-authored-by: Joost Lekkerkerker <[email protected]> * Update config_flow.py --------- Co-authored-by: Joost Lekkerkerker <[email protected]>
home-assistant · Oct 21, 2024 · ca6b759 · ca6b759
1 parent f9d8572
commit ca6b759
Showing 1 changed file with 17 additions and 20 deletions.
diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py
@@ -20,6 +20,7 @@
 
 from homeassistant.components import ssdp
 from homeassistant.config_entries import (
+    SOURCE_REAUTH,
     ConfigEntry,
     ConfigEntryState,
     ConfigFlow,
@@ -86,7 +87,6 @@ def async_get_options_flow(
     def __init__(self) -> None:
         """Initialize the UniFi Network flow."""
         self.config: dict[str, Any] = {}
-        self.reauth_config_entry: ConfigEntry | None = None
         self.reauth_schema: dict[vol.Marker, Any] = {}
 
     async def async_step_user(
@@ -118,13 +118,14 @@ async def async_step_user(
 
             else:
                 if (
-                    self.reauth_config_entry
-                    and self.reauth_config_entry.unique_id is not None
-                    and self.reauth_config_entry.unique_id in self.sites
-                ):
-                    return await self.async_step_site(
-                        {CONF_SITE_ID: self.reauth_config_entry.unique_id}
+                    self.source == SOURCE_REAUTH
+                    and (
+                        (reauth_unique_id := self._get_reauth_entry().unique_id)
+                        is not None
                     )
+                    and reauth_unique_id in self.sites
+                ):
+                    return await self.async_step_site({CONF_SITE_ID: reauth_unique_id})
 
                 return await self.async_step_site()
 
@@ -160,8 +161,8 @@ async def async_step_site(
             config_entry = await self.async_set_unique_id(unique_id)
             abort_reason = "configuration_updated"
 
-            if self.reauth_config_entry:
-                config_entry = self.reauth_config_entry
+            if self.source == SOURCE_REAUTH:
+                config_entry = self._get_reauth_entry()
                 abort_reason = "reauth_successful"
 
             if config_entry:
@@ -192,24 +193,20 @@ async def async_step_reauth(
         self, entry_data: Mapping[str, Any]
     ) -> ConfigFlowResult:
         """Trigger a reauthentication flow."""
-        config_entry = self.hass.config_entries.async_get_entry(
-            self.context["entry_id"]
-        )
-        assert config_entry
-        self.reauth_config_entry = config_entry
+        reauth_entry = self._get_reauth_entry()
 
         self.context["title_placeholders"] = {
-            CONF_HOST: config_entry.data[CONF_HOST],
-            CONF_SITE_ID: config_entry.title,
+            CONF_HOST: reauth_entry.data[CONF_HOST],
+            CONF_SITE_ID: reauth_entry.title,
         }
 
         self.reauth_schema = {
-            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_HOST, default=reauth_entry.data[CONF_HOST]): str,
+            vol.Required(CONF_USERNAME, default=reauth_entry.data[CONF_USERNAME]): str,
             vol.Required(CONF_PASSWORD): str,
-            vol.Required(CONF_PORT, default=config_entry.data[CONF_PORT]): int,
+            vol.Required(CONF_PORT, default=reauth_entry.data[CONF_PORT]): int,
             vol.Required(
-                CONF_VERIFY_SSL, default=config_entry.data[CONF_VERIFY_SSL]
+                CONF_VERIFY_SSL, default=reauth_entry.data[CONF_VERIFY_SSL]
             ): bool,
         }