Fix reauth with esphome when adding noise encryption (#83164)

* Fix reauth with esphome when adding noise encryption

fixes #80813

* fix with unique id

homeassistant/components/esphome/config_flow.py

@@ -17,7 +17,7 @@ from aioesphomeapi import (
 import voluptuous as vol
 
 from homeassistant.components import dhcp, zeroconf
-from homeassistant.config_entries import ConfigFlow
+from homeassistant.config_entries import ConfigEntry, ConfigFlow
 from homeassistant.const import CONF_HOST, CONF_NAME, CONF_PASSWORD, CONF_PORT
 from homeassistant.core import callback
 from homeassistant.data_entry_flow import FlowResult
@@ -40,6 +40,7 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
         self._password: str | None = None
         self._noise_psk: str | None = None
         self._device_info: DeviceInfo | None = None
+        self._reauth_entry: ConfigEntry | None = None
 
     async def _async_step_user_base(
         self, user_input: dict[str, Any] | None = None, error: str | None = None
@@ -72,6 +73,7 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
         """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]
@@ -245,10 +247,11 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
             CONF_PASSWORD: self._password or "",
             CONF_NOISE_PSK: self._noise_psk or "",
         }
-        if "entry_id" in self.context:
+        if self._reauth_entry:
-            entry = self.hass.config_entries.async_get_entry(self.context["entry_id"])
+            entry = self._reauth_entry
-            assert entry is not None
+            self.hass.config_entries.async_update_entry(
-            self.hass.config_entries.async_update_entry(entry, data=config_data)
+                entry, data=self._reauth_entry.data | config_data
+            )
             # Reload the config entry to notify of updated config
             self.hass.async_create_task(
                 self.hass.config_entries.async_reload(entry.entry_id)
@@ -332,6 +335,7 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
 
         self._name = self._device_info.name
         await self.async_set_unique_id(self._name, raise_on_progress=False)
+        if not self._reauth_entry:
             self._abort_if_unique_id_configured(updates={CONF_HOST: self._host})
 
         return None

tests/components/esphome/test_config_flow.py

@@ -559,6 +559,53 @@ async def test_reauth_confirm_invalid(hass, mock_client, mock_zeroconf):
     assert result["errors"]
     assert result["errors"]["base"] == "invalid_psk"
 
+    mock_client.device_info = AsyncMock(return_value=MockDeviceInfo(False, "test"))
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], user_input={CONF_NOISE_PSK: VALID_NOISE_PSK}
+    )
+
+    assert result["type"] == FlowResultType.ABORT
+    assert result["reason"] == "reauth_successful"
+    assert entry.data[CONF_NOISE_PSK] == VALID_NOISE_PSK
+
+
+async def test_reauth_confirm_invalid_with_unique_id(hass, mock_client, mock_zeroconf):
+    """Test reauth initiation with invalid PSK."""
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={CONF_HOST: "127.0.0.1", CONF_PORT: 6053, CONF_PASSWORD: ""},
+        unique_id="test",
+    )
+    entry.add_to_hass(hass)
+
+    result = await hass.config_entries.flow.async_init(
+        "esphome",
+        context={
+            "source": config_entries.SOURCE_REAUTH,
+            "entry_id": entry.entry_id,
+            "unique_id": entry.unique_id,
+        },
+    )
+
+    mock_client.device_info.side_effect = InvalidEncryptionKeyAPIError
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], user_input={CONF_NOISE_PSK: INVALID_NOISE_PSK}
+    )
+
+    assert result["type"] == FlowResultType.FORM
+    assert result["step_id"] == "reauth_confirm"
+    assert result["errors"]
+    assert result["errors"]["base"] == "invalid_psk"
+
+    mock_client.device_info = AsyncMock(return_value=MockDeviceInfo(False, "test"))
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], user_input={CONF_NOISE_PSK: VALID_NOISE_PSK}
+    )
+
+    assert result["type"] == FlowResultType.ABORT
+    assert result["reason"] == "reauth_successful"
+    assert entry.data[CONF_NOISE_PSK] == VALID_NOISE_PSK
+
 
 async def test_discovery_dhcp_updates_host(hass, mock_client):
     """Test dhcp discovery updates host and aborts."""