From 6d72391ee1982ecd1bedac1325b5de8a6be39acd Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Mon, 14 Oct 2024 11:46:14 +0200
Subject: [PATCH] Use reauth helpers in airvisual (#128335)

* Use reauth helpers in airvisual

* Cleanup unused code in tests
---
 homeassistant/components/airvisual/config_flow.py | 14 +++++++++++---
 tests/components/airvisual/test_config_flow.py    |  4 ----
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/homeassistant/components/airvisual/config_flow.py b/homeassistant/components/airvisual/config_flow.py
index 8c012aca93d..7643d541070 100644
--- a/homeassistant/components/airvisual/config_flow.py
+++ b/homeassistant/components/airvisual/config_flow.py
@@ -16,7 +16,12 @@ from pyairvisual.cloud_api import (
 from pyairvisual.errors import AirVisualError
 import voluptuous as vol
 
-from homeassistant.config_entries import ConfigEntry, ConfigFlow, ConfigFlowResult
+from homeassistant.config_entries import (
+    SOURCE_REAUTH,
+    ConfigEntry,
+    ConfigFlow,
+    ConfigFlowResult,
+)
 from homeassistant.const import (
     CONF_API_KEY,
     CONF_COUNTRY,
@@ -140,8 +145,11 @@ class AirVisualFlowHandler(ConfigFlow, domain=DOMAIN):
 
                 valid_keys.add(user_input[CONF_API_KEY])
 
-        if existing_entry := await self.async_set_unique_id(self._geo_id):
-            return self.async_update_reload_and_abort(existing_entry, data=user_input)
+        if self.source == SOURCE_REAUTH:
+            return self.async_update_reload_and_abort(
+                self._get_reauth_entry(),
+                data_updates={CONF_API_KEY: user_input[CONF_API_KEY]},
+            )
 
         return self.async_create_entry(
             title=f"Cloud API ({self._geo_id})",
diff --git a/tests/components/airvisual/test_config_flow.py b/tests/components/airvisual/test_config_flow.py
index e38fc64587e..632bdb72eb4 100644
--- a/tests/components/airvisual/test_config_flow.py
+++ b/tests/components/airvisual/test_config_flow.py
@@ -155,10 +155,6 @@ async def test_step_reauth(
     assert result["type"] is FlowResultType.FORM
     assert result["step_id"] == "reauth_confirm"
 
-    result = await hass.config_entries.flow.async_configure(result["flow_id"])
-    assert result["type"] is FlowResultType.FORM
-    assert result["step_id"] == "reauth_confirm"
-
     new_api_key = "defgh67890"
 
     result = await hass.config_entries.flow.async_configure(