diff --git a/homeassistant/components/ambee/__init__.py b/homeassistant/components/ambee/__init__.py index 4968420174e..362dc26d851 100644 --- a/homeassistant/components/ambee/__init__.py +++ b/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}: - coordinator: DataUpdateCoordinator = DataUpdateCoordinator( - hass, - LOGGER, - name=DOMAIN, - update_interval=SCAN_INTERVAL, - update_method=getattr(client, service), - ) - await coordinator.async_config_entry_first_refresh() - hass.data[DOMAIN][entry.entry_id][service] = coordinator + async def update_air_quality() -> AirQuality: + """Update method for updating Ambee Air Quality data.""" + try: + return await client.air_quality() + except AmbeeAuthenticationError as err: + raise ConfigEntryAuthFailed from err + + air_quality: DataUpdateCoordinator[AirQuality] = DataUpdateCoordinator( + hass, + 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 diff --git a/homeassistant/components/ambee/config_flow.py b/homeassistant/components/ambee/config_flow.py index 78b25e6226c..0550c541ed0 100644 --- a/homeassistant/components/ambee/config_flow.py +++ b/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, + ) diff --git a/homeassistant/components/ambee/strings.json b/homeassistant/components/ambee/strings.json index 8bec71ebe29..e3c306788dd 100644 --- a/homeassistant/components/ambee/strings.json +++ b/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%]" } } } diff --git a/homeassistant/components/ambee/translations/en.json b/homeassistant/components/ambee/translations/en.json index 8728f56fb4e..433580e8023 100644 --- a/homeassistant/components/ambee/translations/en.json +++ b/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", diff --git a/tests/components/ambee/test_config_flow.py b/tests/components/ambee/test_config_flow.py index 1f91a842321..a6220418681 100644 --- a/tests/components/ambee/test_config_flow.py +++ b/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"} diff --git a/tests/components/ambee/test_init.py b/tests/components/ambee/test_init.py index 5db3255dc53..c6ad45735ff 100644 --- a/tests/components/ambee/test_init.py +++ b/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