Add reauth flow to co2signal (home-assistant#104507)

homeassistant/components/co2signal/config_flow.py

@@ -1,13 +1,15 @@
 """Config flow for Co2signal integration."""
 from __future__ import annotations
 
+from collections.abc import Mapping
 from typing import Any
 
 from aioelectricitymaps import ElectricityMaps
 from aioelectricitymaps.exceptions import ElectricityMapsError, InvalidToken
 import voluptuous as vol
 
 from homeassistant import config_entries
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE
 from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
@@ -32,6 +34,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
     VERSION = 1
     _data: dict | None
+    _reauth_entry: ConfigEntry | None = None
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
@@ -113,25 +116,52 @@ async def async_step_country(
             "country", data_schema, {**self._data, **user_input}
         )
 
+    async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult:
+        """Handle the reauth step."""
+        self._reauth_entry = self.hass.config_entries.async_get_entry(
+            self.context["entry_id"]
+        )
+
+        data_schema = vol.Schema(
+            {
+                vol.Required(CONF_API_KEY): cv.string,
+            }
+        )
+        return await self._validate_and_create("reauth", data_schema, entry_data)
+
     async def _validate_and_create(
-        self, step_id: str, data_schema: vol.Schema, data: dict
+        self, step_id: str, data_schema: vol.Schema, data: Mapping[str, Any]
     ) -> FlowResult:
         """Validate data and show form if it is invalid."""
         errors: dict[str, str] = {}
 
-        session = async_get_clientsession(self.hass)
-        em = ElectricityMaps(token=data[CONF_API_KEY], session=session)
-        try:
-            await fetch_latest_carbon_intensity(self.hass, em, data)
-        except InvalidToken:
-            errors["base"] = "invalid_auth"
-        except ElectricityMapsError:
-            errors["base"] = "unknown"
-        else:
-            return self.async_create_entry(
-                title=get_extra_name(data) or "CO2 Signal",
-                data=data,
-            )
+        if data:
+            session = async_get_clientsession(self.hass)
+            em = ElectricityMaps(token=data[CONF_API_KEY], session=session)
+
+            try:
+                await fetch_latest_carbon_intensity(self.hass, em, data)
+            except InvalidToken:
+                errors["base"] = "invalid_auth"
+            except ElectricityMapsError:
+                errors["base"] = "unknown"
+            else:
+                if self._reauth_entry:
+                    self.hass.config_entries.async_update_entry(
+                        self._reauth_entry,
+                        data={
+                            CONF_API_KEY: data[CONF_API_KEY],
+                        },
+                    )
+                    await self.hass.config_entries.async_reload(
+                        self._reauth_entry.entry_id
+                    )
+                    return self.async_abort(reason="reauth_successful")
+
+                return self.async_create_entry(
+                    title=get_extra_name(data) or "CO2 Signal",
+                    data=data,
+                )
 
         return self.async_show_form(
             step_id=step_id,

homeassistant/components/co2signal/coordinator.py

@@ -10,7 +10,7 @@
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryError
+from homeassistant.exceptions import ConfigEntryAuthFailed
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import DOMAIN
@@ -44,6 +44,6 @@ async def _async_update_data(self) -> CarbonIntensityResponse:
                 self.hass, self.client, self.config_entry.data
             )
         except InvalidToken as err:
-            raise ConfigEntryError from err
+            raise ConfigEntryAuthFailed from err
         except ElectricityMapsError as err:
             raise UpdateFailed(str(err)) from err

homeassistant/components/co2signal/helpers.py

@@ -1,5 +1,5 @@
 """Helper functions for the CO2 Signal integration."""
-from types import MappingProxyType
+from collections.abc import Mapping
 from typing import Any
 
 from aioelectricitymaps import ElectricityMaps
@@ -14,7 +14,7 @@
 async def fetch_latest_carbon_intensity(
     hass: HomeAssistant,
     em: ElectricityMaps,
-    config: dict[str, Any] | MappingProxyType[str, Any],
+    config: Mapping[str, Any],
 ) -> CarbonIntensityResponse:
     """Fetch the latest carbon intensity based on country code or location coordinates."""
     if CONF_COUNTRY_CODE in config:

homeassistant/components/co2signal/strings.json

@@ -18,6 +18,11 @@
         "data": {
           "country_code": "Country code"
         }
+      },
+      "reauth": {
+        "data": {
+          "api_key": "[%key:common::config_flow::data::access_token%]"
+        }
       }
     },
     "error": {
@@ -28,7 +33,8 @@
     "abort": {
       "already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
       "unknown": "[%key:common::config_flow::error::unknown%]",
-      "api_ratelimit": "[%key:component::co2signal::config::error::api_ratelimit%]"
+      "api_ratelimit": "[%key:component::co2signal::config::error::api_ratelimit%]",
+      "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]"
     }
   },
   "entity": {

tests/components/co2signal/test_config_flow.py

@@ -10,9 +10,12 @@
 
 from homeassistant import config_entries
 from homeassistant.components.co2signal import DOMAIN, config_flow
+from homeassistant.const import CONF_API_KEY
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 
+from tests.common import MockConfigEntry
+
 
 @pytest.mark.usefixtures("electricity_maps")
 async def test_form_home(hass: HomeAssistant) -> None:
@@ -186,3 +189,40 @@ async def test_form_error_handling(
     assert result["data"] == {
         "api_key": "api_key",
     }
+
+
+async def test_reauth(
+    hass: HomeAssistant,
+    config_entry: MockConfigEntry,
+    electricity_maps: AsyncMock,
+) -> None:
+    """Test reauth flow."""
+    config_entry.add_to_hass(hass)
+
+    init_result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={
+            "source": config_entries.SOURCE_REAUTH,
+            "entry_id": config_entry.entry_id,
+        },
+        data=None,
+    )
+
+    assert init_result["type"] == FlowResultType.FORM
+    assert init_result["step_id"] == "reauth"
+
+    with patch(
+        "homeassistant.components.co2signal.async_setup_entry",
+        return_value=True,
+    ) as mock_setup_entry:
+        configure_result = await hass.config_entries.flow.async_configure(
+            init_result["flow_id"],
+            {
+                CONF_API_KEY: "api_key2",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert configure_result["type"] == FlowResultType.ABORT
+    assert configure_result["reason"] == "reauth_successful"
+    assert len(mock_setup_entry.mock_calls) == 1

tests/components/co2signal/test_sensor.py

@@ -42,7 +42,6 @@ async def test_sensor(
 @pytest.mark.parametrize(
     "error",
     [
-        InvalidToken,
         ElectricityMapsDecodeError,
         ElectricityMapsError,
         Exception,
@@ -82,3 +81,25 @@ async def test_sensor_update_fail(
     assert (state := hass.states.get("sensor.electricity_maps_co2_intensity"))
     assert state.state == "45.9862319009581"
     assert len(electricity_maps.mock_calls) == 3
+
+
+@pytest.mark.usefixtures("setup_integration")
+async def test_sensor_reauth_triggered(
+    hass: HomeAssistant,
+    freezer: FrozenDateTimeFactory,
+    electricity_maps: AsyncMock,
+):
+    """Test if reauth flow is triggered."""
+    assert (state := hass.states.get("sensor.electricity_maps_co2_intensity"))
+    assert state.state == "45.9862319009581"
+
+    electricity_maps.latest_carbon_intensity_by_coordinates.side_effect = InvalidToken
+    electricity_maps.latest_carbon_intensity_by_country_code.side_effect = InvalidToken
+
+    freezer.tick(timedelta(minutes=20))
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    assert (flows := hass.config_entries.flow.async_progress())
+    assert len(flows) == 1
+    assert flows[0]["step_id"] == "reauth"