Add re-authentication support to Ambee (#51773)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/ambee/__init__.py

@@ -1,12 +1,13 @@
 """Support for Ambee."""
 from __future__ import annotations
 
-from ambee import Ambee
+from ambee import AirQuality, Ambee, AmbeeAuthenticationError, Pollen
 
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryAuthFailed
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from .const import DOMAIN, LOGGER, SCAN_INTERVAL, SERVICE_AIR_QUALITY, SERVICE_POLLEN
@@ -24,16 +25,39 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         longitude=entry.data[CONF_LONGITUDE],
     )
 
-    for service in {SERVICE_AIR_QUALITY, SERVICE_POLLEN}:
+    async def update_air_quality() -> AirQuality:
-        coordinator: DataUpdateCoordinator = DataUpdateCoordinator(
+        """Update method for updating Ambee Air Quality data."""
-            hass,
+        try:
-            LOGGER,
+            return await client.air_quality()
-            name=DOMAIN,
+        except AmbeeAuthenticationError as err:
-            update_interval=SCAN_INTERVAL,
+            raise ConfigEntryAuthFailed from err
-            update_method=getattr(client, service),
+
-        )
+    air_quality: DataUpdateCoordinator[AirQuality] = DataUpdateCoordinator(
-        await coordinator.async_config_entry_first_refresh()
+        hass,
-        hass.data[DOMAIN][entry.entry_id][service] = coordinator
+        LOGGER,
+        name=f"{DOMAIN}_{SERVICE_AIR_QUALITY}",
+        update_interval=SCAN_INTERVAL,
+        update_method=update_air_quality,
+    )
+    await air_quality.async_config_entry_first_refresh()
+    hass.data[DOMAIN][entry.entry_id][SERVICE_AIR_QUALITY] = air_quality
+
+    async def update_pollen() -> Pollen:
+        """Update method for updating Ambee Pollen data."""
+        try:
+            return await client.pollen()
+        except AmbeeAuthenticationError as err:
+            raise ConfigEntryAuthFailed from err
+
+    pollen: DataUpdateCoordinator[Pollen] = DataUpdateCoordinator(
+        hass,
+        LOGGER,
+        name=f"{DOMAIN}_{SERVICE_POLLEN}",
+        update_interval=SCAN_INTERVAL,
+        update_method=update_pollen,
+    )
+    await pollen.async_config_entry_first_refresh()
+    hass.data[DOMAIN][entry.entry_id][SERVICE_POLLEN] = pollen
 
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
     return True

homeassistant/components/ambee/config_flow.py

@@ -6,7 +6,7 @@ from typing import Any
 from ambee import Ambee, AmbeeAuthenticationError, AmbeeError
 import voluptuous as vol
 
-from homeassistant.config_entries import ConfigFlow
+from homeassistant.config_entries import ConfigEntry, ConfigFlow
 from homeassistant.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE, CONF_NAME
 from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
@@ -20,6 +20,8 @@ class AmbeeFlowHandler(ConfigFlow, domain=DOMAIN):
 
     VERSION = 1
 
+    entry: ConfigEntry | None = None
+
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
     ) -> FlowResult:
@@ -68,3 +70,46 @@ class AmbeeFlowHandler(ConfigFlow, domain=DOMAIN):
             ),
             errors=errors,
         )
+
+    async def async_step_reauth(self, data: dict[str, Any]) -> FlowResult:
+        """Handle initiation of re-authentication with Ambee."""
+        self.entry = self.hass.config_entries.async_get_entry(self.context["entry_id"])
+        return await self.async_step_reauth_confirm()
+
+    async def async_step_reauth_confirm(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Handle re-authentication with Ambee."""
+        errors = {}
+        if user_input is not None and self.entry:
+            session = async_get_clientsession(self.hass)
+            client = Ambee(
+                api_key=user_input[CONF_API_KEY],
+                latitude=self.entry.data[CONF_LATITUDE],
+                longitude=self.entry.data[CONF_LONGITUDE],
+                session=session,
+            )
+            try:
+                await client.air_quality()
+            except AmbeeAuthenticationError:
+                errors["base"] = "invalid_api_key"
+            except AmbeeError:
+                errors["base"] = "cannot_connect"
+            else:
+                self.hass.config_entries.async_update_entry(
+                    self.entry,
+                    data={
+                        **self.entry.data,
+                        CONF_API_KEY: user_input[CONF_API_KEY],
+                    },
+                )
+                self.hass.async_create_task(
+                    self.hass.config_entries.async_reload(self.entry.entry_id)
+                )
+                return self.async_abort(reason="reauth_successful")
+
+        return self.async_show_form(
+            step_id="reauth_confirm",
+            data_schema=vol.Schema({vol.Required(CONF_API_KEY): str}),
+            errors=errors,
+        )

homeassistant/components/ambee/strings.json

@@ -9,11 +9,20 @@
           "longitude": "[%key:common::config_flow::data::longitude%]",
           "name": "[%key:common::config_flow::data::name%]"
         }
+      },
+      "reauth_confirm": {
+        "data": {
+          "description": "Re-authenticate with your Ambee account.",
+          "api_key": "[%key:common::config_flow::data::api_key%]"
+        }
       }
     },
     "error": {
       "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
       "invalid_api_key": "[%key:common::config_flow::error::invalid_api_key%]"
+    },
+    "abort": {
+      "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]"
     }
   }
 }

homeassistant/components/ambee/translations/en.json

@@ -1,10 +1,19 @@
 {
     "config": {
+        "abort": {
+            "reauth_successful": "Re-authentication was successful"
+        },
         "error": {
             "cannot_connect": "Failed to connect",
             "invalid_api_key": "Invalid API key"
         },
         "step": {
+            "reauth_confirm": {
+                "data": {
+                    "api_key": "API Key",
+                    "description": "Re-authenticate with your Ambee account."
+                }
+            },
             "user": {
                 "data": {
                     "api_key": "API Key",

tests/components/ambee/test_config_flow.py

@@ -5,10 +5,16 @@ from unittest.mock import patch
 from ambee import AmbeeAuthenticationError, AmbeeError
 
 from homeassistant.components.ambee.const import DOMAIN
-from homeassistant.config_entries import SOURCE_USER
+from homeassistant.config_entries import SOURCE_REAUTH, SOURCE_USER
 from homeassistant.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE, CONF_NAME
 from homeassistant.core import HomeAssistant
-from homeassistant.data_entry_flow import RESULT_TYPE_CREATE_ENTRY, RESULT_TYPE_FORM
+from homeassistant.data_entry_flow import (
+    RESULT_TYPE_ABORT,
+    RESULT_TYPE_CREATE_ENTRY,
+    RESULT_TYPE_FORM,
+)
+
+from tests.common import MockConfigEntry
 
 
 async def test_full_user_flow(hass: HomeAssistant) -> None:
@@ -127,3 +133,140 @@ async def test_api_error(hass: HomeAssistant) -> None:
 
     assert result.get("type") == RESULT_TYPE_FORM
     assert result.get("errors") == {"base": "cannot_connect"}
+
+
+async def test_reauth_flow(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
+    """Test the reauthentication configuration flow."""
+    mock_config_entry.add_to_hass(hass)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={
+            "source": SOURCE_REAUTH,
+            "unique_id": mock_config_entry.unique_id,
+            "entry_id": mock_config_entry.entry_id,
+        },
+        data=mock_config_entry.data,
+    )
+    assert result.get("type") == RESULT_TYPE_FORM
+    assert result.get("step_id") == "reauth_confirm"
+    assert "flow_id" in result
+
+    with patch(
+        "homeassistant.components.ambee.config_flow.Ambee.air_quality"
+    ) as mock_ambee, patch(
+        "homeassistant.components.ambee.async_setup_entry", return_value=True
+    ) as mock_setup_entry:
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {CONF_API_KEY: "other_key"},
+        )
+        await hass.async_block_till_done()
+
+    assert result2.get("type") == RESULT_TYPE_ABORT
+    assert result2.get("reason") == "reauth_successful"
+    assert mock_config_entry.data == {
+        CONF_API_KEY: "other_key",
+        CONF_LATITUDE: 52.42,
+        CONF_LONGITUDE: 4.44,
+    }
+
+    assert len(mock_ambee.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 1
+
+
+async def test_reauth_with_authentication_error(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
+    """Test the reauthentication configuration flow with an authentication error.
+
+    This tests tests a reauth flow, with a case the user enters an invalid
+    API token, but recover by entering the correct one.
+    """
+    mock_config_entry.add_to_hass(hass)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={
+            "source": SOURCE_REAUTH,
+            "unique_id": mock_config_entry.unique_id,
+            "entry_id": mock_config_entry.entry_id,
+        },
+        data=mock_config_entry.data,
+    )
+    assert result.get("type") == RESULT_TYPE_FORM
+    assert result.get("step_id") == "reauth_confirm"
+    assert "flow_id" in result
+
+    with patch(
+        "homeassistant.components.ambee.config_flow.Ambee.air_quality",
+        side_effect=AmbeeAuthenticationError,
+    ):
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            user_input={
+                CONF_API_KEY: "invalid",
+            },
+        )
+
+    assert result2.get("type") == RESULT_TYPE_FORM
+    assert result2.get("step_id") == "reauth_confirm"
+    assert result2.get("errors") == {"base": "invalid_api_key"}
+    assert "flow_id" in result2
+
+    with patch(
+        "homeassistant.components.ambee.config_flow.Ambee.air_quality"
+    ) as mock_ambee, patch(
+        "homeassistant.components.ambee.async_setup_entry", return_value=True
+    ) as mock_setup_entry:
+        result3 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {CONF_API_KEY: "other_key"},
+        )
+        await hass.async_block_till_done()
+
+    assert result3.get("type") == RESULT_TYPE_ABORT
+    assert result3.get("reason") == "reauth_successful"
+    assert mock_config_entry.data == {
+        CONF_API_KEY: "other_key",
+        CONF_LATITUDE: 52.42,
+        CONF_LONGITUDE: 4.44,
+    }
+
+    assert len(mock_ambee.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 1
+
+
+async def test_reauth_api_error(
+    hass: HomeAssistant, mock_config_entry: MockConfigEntry
+) -> None:
+    """Test API error during reauthentication."""
+    mock_config_entry.add_to_hass(hass)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={
+            "source": SOURCE_REAUTH,
+            "unique_id": mock_config_entry.unique_id,
+            "entry_id": mock_config_entry.entry_id,
+        },
+        data=mock_config_entry.data,
+    )
+    assert "flow_id" in result
+
+    with patch(
+        "homeassistant.components.ambee.config_flow.Ambee.air_quality",
+        side_effect=AmbeeError,
+    ):
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            user_input={
+                CONF_API_KEY: "invalid",
+            },
+        )
+
+    assert result2.get("type") == RESULT_TYPE_FORM
+    assert result2.get("step_id") == "reauth_confirm"
+    assert result2.get("errors") == {"base": "cannot_connect"}

tests/components/ambee/test_init.py

@@ -2,9 +2,11 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 from ambee import AmbeeConnectionError
+from ambee.exceptions import AmbeeAuthenticationError
+import pytest
 
 from homeassistant.components.ambee.const import DOMAIN
-from homeassistant.config_entries import ConfigEntryState
+from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState
 from homeassistant.core import HomeAssistant
 
 from tests.common import MockConfigEntry
@@ -44,3 +46,33 @@ async def test_config_entry_not_ready(
 
     assert mock_request.call_count == 1
     assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY
+
+
+@pytest.mark.parametrize("service_name", ["air_quality", "pollen"])
+async def test_config_entry_authentication_failed(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_ambee: MagicMock,
+    service_name: str,
+) -> None:
+    """Test the Ambee configuration entry not ready."""
+    mock_config_entry.add_to_hass(hass)
+
+    service = getattr(mock_ambee.return_value, service_name)
+    service.side_effect = AmbeeAuthenticationError
+
+    await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert mock_config_entry.state is ConfigEntryState.SETUP_ERROR
+
+    flows = hass.config_entries.flow.async_progress()
+    assert len(flows) == 1
+
+    flow = flows[0]
+    assert flow.get("step_id") == "reauth_confirm"
+    assert flow.get("handler") == DOMAIN
+
+    assert "context" in flow
+    assert flow["context"].get("source") == SOURCE_REAUTH
+    assert flow["context"].get("entry_id") == mock_config_entry.entry_id