From a218f4adc3f12e2cf0c84028ab683c4a28369134 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Thu, 3 Oct 2024 14:38:14 +0200
Subject: [PATCH] Use reauth helpers in esphome config flow (#127419)

Use _get_reauth_entry in esphome config flow
---
 .../components/esphome/config_flow.py         | 22 +++++++++----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/homeassistant/components/esphome/config_flow.py b/homeassistant/components/esphome/config_flow.py
index 59f70d6c6b6..937cad040ea 100644
--- a/homeassistant/components/esphome/config_flow.py
+++ b/homeassistant/components/esphome/config_flow.py
@@ -23,6 +23,7 @@ import voluptuous as vol
 from homeassistant.components import dhcp, zeroconf
 from homeassistant.components.hassio import HassioServiceInfo
 from homeassistant.config_entries import (
+    SOURCE_REAUTH,
     ConfigEntry,
     ConfigFlow,
     ConfigFlowResult,
@@ -57,6 +58,8 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
 
     VERSION = 1
 
+    _reauth_entry: ConfigEntry
+
     def __init__(self) -> None:
         """Initialize flow."""
         self._host: str | None = None
@@ -66,7 +69,6 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
         self._noise_required: bool | None = None
         self._noise_psk: str | None = None
         self._device_info: DeviceInfo | None = None
-        self._reauth_entry: ConfigEntry | None = None
         # The ESPHome name as per its config
         self._device_name: str | None = None
 
@@ -103,14 +105,12 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
         self, entry_data: Mapping[str, Any]
     ) -> ConfigFlowResult:
         """Handle a flow initialized by a reauth event."""
-        entry = self.hass.config_entries.async_get_entry(self.context["entry_id"])
-        assert entry is not None
-        self._reauth_entry = entry
-        self._host = entry.data[CONF_HOST]
-        self._port = entry.data[CONF_PORT]
-        self._password = entry.data[CONF_PASSWORD]
-        self._name = entry.title
-        self._device_name = entry.data.get(CONF_DEVICE_NAME)
+        self._reauth_entry = self._get_reauth_entry()
+        self._host = entry_data[CONF_HOST]
+        self._port = entry_data[CONF_PORT]
+        self._password = entry_data[CONF_PASSWORD]
+        self._name = self._reauth_entry.title
+        self._device_name = entry_data.get(CONF_DEVICE_NAME)
 
         # Device without encryption allows fetching device info. We can then check
         # if the device is no longer using a password. If we did try with a password,
@@ -324,7 +324,7 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
         config_options = {
             CONF_ALLOW_SERVICE_CALLS: DEFAULT_NEW_CONFIG_ALLOW_ALLOW_SERVICE_CALLS,
         }
-        if self._reauth_entry:
+        if self.source == SOURCE_REAUTH:
             return self.async_update_reload_and_abort(
                 self._reauth_entry, data=self._reauth_entry.data | config_data
             )
@@ -411,7 +411,7 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
         self._device_name = self._device_info.name
         mac_address = format_mac(self._device_info.mac_address)
         await self.async_set_unique_id(mac_address, raise_on_progress=False)
-        if not self._reauth_entry:
+        if self.source != SOURCE_REAUTH:
             self._abort_if_unique_id_configured(
                 updates={CONF_HOST: self._host, CONF_PORT: self._port}
             )