Add re-authentication support to Ambee (#51773)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
This commit is contained in:
Franck Nijhof 2021-06-12 16:18:06 +02:00 committed by GitHub
parent cfce71d7df
commit c242e56b8c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 277 additions and 15 deletions

View File

@ -1,12 +1,13 @@
"""Support for Ambee.""" """Support for Ambee."""
from __future__ import annotations 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.components.sensor import DOMAIN as SENSOR_DOMAIN
from homeassistant.config_entries import ConfigEntry from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE from homeassistant.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryAuthFailed
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
from .const import DOMAIN, LOGGER, SCAN_INTERVAL, SERVICE_AIR_QUALITY, SERVICE_POLLEN 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], 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) hass.config_entries.async_setup_platforms(entry, PLATFORMS)
return True return True

View File

@ -6,7 +6,7 @@ from typing import Any
from ambee import Ambee, AmbeeAuthenticationError, AmbeeError from ambee import Ambee, AmbeeAuthenticationError, AmbeeError
import voluptuous as vol 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.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE, CONF_NAME
from homeassistant.data_entry_flow import FlowResult from homeassistant.data_entry_flow import FlowResult
from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.aiohttp_client import async_get_clientsession
@ -20,6 +20,8 @@ class AmbeeFlowHandler(ConfigFlow, domain=DOMAIN):
VERSION = 1 VERSION = 1
entry: ConfigEntry | None = None
async def async_step_user( async def async_step_user(
self, user_input: dict[str, Any] | None = None self, user_input: dict[str, Any] | None = None
) -> FlowResult: ) -> FlowResult:
@ -68,3 +70,46 @@ class AmbeeFlowHandler(ConfigFlow, domain=DOMAIN):
), ),
errors=errors, 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,
)

View File

@ -9,11 +9,20 @@
"longitude": "[%key:common::config_flow::data::longitude%]", "longitude": "[%key:common::config_flow::data::longitude%]",
"name": "[%key:common::config_flow::data::name%]" "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": { "error": {
"cannot_connect": "[%key:common::config_flow::error::cannot_connect%]", "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
"invalid_api_key": "[%key:common::config_flow::error::invalid_api_key%]" "invalid_api_key": "[%key:common::config_flow::error::invalid_api_key%]"
},
"abort": {
"reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]"
} }
} }
} }

View File

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

View File

@ -5,10 +5,16 @@ from unittest.mock import patch
from ambee import AmbeeAuthenticationError, AmbeeError from ambee import AmbeeAuthenticationError, AmbeeError
from homeassistant.components.ambee.const import DOMAIN 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.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE, CONF_NAME
from homeassistant.core import HomeAssistant 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: 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("type") == RESULT_TYPE_FORM
assert result.get("errors") == {"base": "cannot_connect"} 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"}

View File

@ -2,9 +2,11 @@
from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import AsyncMock, MagicMock, patch
from ambee import AmbeeConnectionError from ambee import AmbeeConnectionError
from ambee.exceptions import AmbeeAuthenticationError
import pytest
from homeassistant.components.ambee.const import DOMAIN 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 homeassistant.core import HomeAssistant
from tests.common import MockConfigEntry from tests.common import MockConfigEntry
@ -44,3 +46,33 @@ async def test_config_entry_not_ready(
assert mock_request.call_count == 1 assert mock_request.call_count == 1
assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY 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