Use reauth helpers in google_generative_ai_conversation (home-assista…

…nt#128583)
aschmere · Oct 17, 2024 · 536d702 · 536d702
1 parent 420070a
commit 536d702
Showing 1 changed file with 7 additions and 12 deletions.
diff --git a/homeassistant/components/google_generative_ai_conversation/config_flow.py b/homeassistant/components/google_generative_ai_conversation/config_flow.py
@@ -15,6 +15,7 @@
 import voluptuous as vol
 
 from homeassistant.config_entries import (
+    SOURCE_REAUTH,
     ConfigEntry,
     ConfigFlow,
     ConfigFlowResult,
@@ -85,10 +86,6 @@ class GoogleGenerativeAIConfigFlow(ConfigFlow, domain=DOMAIN):
 
     VERSION = 1
 
-    def __init__(self) -> None:
-        """Initialize a new GoogleGenerativeAIConfigFlow."""
-        self.reauth_entry: ConfigEntry | None = None
-
     async def async_step_api(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
@@ -106,9 +103,9 @@ async def async_step_api(
                 _LOGGER.exception("Unexpected exception")
                 errors["base"] = "unknown"
             else:
-                if self.reauth_entry:
+                if self.source == SOURCE_REAUTH:
                     return self.async_update_reload_and_abort(
-                        self.reauth_entry,
+                        self._get_reauth_entry(),
                         data=user_input,
                     )
                 return self.async_create_entry(
@@ -135,9 +132,6 @@ async def async_step_reauth(
         self, entry_data: Mapping[str, Any]
     ) -> ConfigFlowResult:
         """Handle configuration by re-auth."""
-        self.reauth_entry = self.hass.config_entries.async_get_entry(
-            self.context["entry_id"]
-        )
         return await self.async_step_reauth_confirm()
 
     async def async_step_reauth_confirm(
@@ -146,12 +140,13 @@ async def async_step_reauth_confirm(
         """Dialog that informs the user that reauth is required."""
         if user_input is not None:
             return await self.async_step_api()
-        assert self.reauth_entry
+
+        reauth_entry = self._get_reauth_entry()
         return self.async_show_form(
             step_id="reauth_confirm",
             description_placeholders={
-                CONF_NAME: self.reauth_entry.title,
-                CONF_API_KEY: self.reauth_entry.data.get(CONF_API_KEY, ""),
+                CONF_NAME: reauth_entry.title,
+                CONF_API_KEY: reauth_entry.data.get(CONF_API_KEY, ""),
             },
         )