From 3f7c6055d45999de42c00b24db9e10308029aefd Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Wed, 2 Oct 2024 14:33:29 +0200
Subject: [PATCH] Use _get_reauth/reconfigure_entry in lamarzocco (#127298)

---
 .../components/lamarzocco/config_flow.py      | 28 ++++++++-----------
 1 file changed, 11 insertions(+), 17 deletions(-)

diff --git a/homeassistant/components/lamarzocco/config_flow.py b/homeassistant/components/lamarzocco/config_flow.py
index 5a5cad00f64..92d428ffebe 100644
--- a/homeassistant/components/lamarzocco/config_flow.py
+++ b/homeassistant/components/lamarzocco/config_flow.py
@@ -15,6 +15,8 @@ from homeassistant.components.bluetooth import (
     async_discovered_service_info,
 )
 from homeassistant.config_entries import (
+    SOURCE_REAUTH,
+    SOURCE_RECONFIGURE,
     ConfigEntry,
     ConfigFlow,
     ConfigFlowResult,
@@ -52,11 +54,11 @@ class LmConfigFlow(ConfigFlow, domain=DOMAIN):
 
     VERSION = 2
 
+    reauth_entry: ConfigEntry
+    reconfigure_entry: ConfigEntry
+
     def __init__(self) -> None:
         """Initialize the config flow."""
-
-        self.reauth_entry: ConfigEntry | None = None
-        self.reconfigure_entry: ConfigEntry | None = None
         self._config: dict[str, Any] = {}
         self._fleet: dict[str, LaMarzoccoDeviceInfo] = {}
         self._discovered: dict[str, str] = {}
@@ -70,7 +72,7 @@ class LmConfigFlow(ConfigFlow, domain=DOMAIN):
 
         if user_input:
             data: dict[str, Any] = {}
-            if self.reauth_entry:
+            if self.source == SOURCE_REAUTH:
                 data = dict(self.reauth_entry.data)
             data = {
                 **data,
@@ -95,7 +97,7 @@ class LmConfigFlow(ConfigFlow, domain=DOMAIN):
                     errors["base"] = "no_machines"
 
             if not errors:
-                if self.reauth_entry:
+                if self.source == SOURCE_REAUTH:
                     return self.async_update_reload_and_abort(
                         self.reauth_entry, data=data, reason="reauth_successful"
                     )
@@ -134,7 +136,7 @@ class LmConfigFlow(ConfigFlow, domain=DOMAIN):
         if user_input:
             if not self._discovered:
                 serial_number = user_input[CONF_MACHINE]
-                if self.reconfigure_entry is None:
+                if self.source != SOURCE_RECONFIGURE:
                     await self.async_set_unique_id(serial_number)
                     self._abort_if_unique_id_configured()
             else:
@@ -154,7 +156,7 @@ class LmConfigFlow(ConfigFlow, domain=DOMAIN):
                     self._config[CONF_HOST] = user_input[CONF_HOST]
 
             if not errors:
-                if self.reconfigure_entry:
+                if self.source == SOURCE_RECONFIGURE:
                     for service_info in async_discovered_service_info(self.hass):
                         self._discovered[service_info.name] = service_info.address
 
@@ -204,8 +206,6 @@ class LmConfigFlow(ConfigFlow, domain=DOMAIN):
     ) -> ConfigFlowResult:
         """Handle Bluetooth device selection."""
 
-        assert self.reconfigure_entry
-
         if user_input is not None:
             return self.async_update_reload_and_abort(
                 self.reconfigure_entry,
@@ -266,9 +266,7 @@ class LmConfigFlow(ConfigFlow, domain=DOMAIN):
         self, entry_data: Mapping[str, Any]
     ) -> ConfigFlowResult:
         """Perform reauth upon an API authentication error."""
-        self.reauth_entry = self.hass.config_entries.async_get_entry(
-            self.context["entry_id"]
-        )
+        self.reauth_entry = self._get_reauth_entry()
         return await self.async_step_reauth_confirm()
 
     async def async_step_reauth_confirm(
@@ -291,17 +289,13 @@ class LmConfigFlow(ConfigFlow, domain=DOMAIN):
         self, entry_data: Mapping[str, Any]
     ) -> ConfigFlowResult:
         """Perform reconfiguration of the config entry."""
-        self.reconfigure_entry = self.hass.config_entries.async_get_entry(
-            self.context["entry_id"]
-        )
+        self.reconfigure_entry = self._get_reconfigure_entry()
         return await self.async_step_reconfigure_confirm()
 
     async def async_step_reconfigure_confirm(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Confirm reconfiguration of the device."""
-        assert self.reconfigure_entry
-
         if not user_input:
             return self.async_show_form(
                 step_id="reconfigure_confirm",