From 1cace9a609d20a02887208311a500bc145070cf3 Mon Sep 17 00:00:00 2001 From: Brett Adams Date: Mon, 8 Apr 2024 17:44:51 +1000 Subject: [PATCH] Add reauth to Teslemetry (#114726) * Add reauth * Add tests * PARALLEL_UPDATES * Bump quality to platinum * Fix assertion * Remove quality * Remove async_create_task * Review Feedback * Remove loop inside parametrize * Change config during reauth * Fix missing return --- .../components/teslemetry/__init__.py | 14 ++- .../components/teslemetry/config_flow.py | 75 +++++++++++---- .../components/teslemetry/coordinator.py | 21 +++- .../components/teslemetry/test_config_flow.py | 96 +++++++++++++++++++ tests/components/teslemetry/test_init.py | 61 ++++++------ 5 files changed, 207 insertions(+), 60 deletions(-) diff --git a/homeassistant/components/teslemetry/__init__.py b/homeassistant/components/teslemetry/__init__.py index 1da3533fef1..084d51ff31b 100644 --- a/homeassistant/components/teslemetry/__init__.py +++ b/homeassistant/components/teslemetry/__init__.py @@ -13,10 +13,10 @@ from tesla_fleet_api.exceptions import ( from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_ACCESS_TOKEN, Platform from homeassistant.core import HomeAssistant -from homeassistant.exceptions import ConfigEntryNotReady +from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady from homeassistant.helpers.aiohttp_client import async_get_clientsession -from .const import DOMAIN, LOGGER +from .const import DOMAIN from .coordinator import ( TeslemetryEnergyDataCoordinator, TeslemetryVehicleDataCoordinator, @@ -38,12 +38,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: ) try: products = (await teslemetry.products())["response"] - except InvalidToken: - LOGGER.error("Access token is invalid, unable to connect to Teslemetry") - return False - except SubscriptionRequired: - LOGGER.error("Subscription required, unable to connect to Telemetry") - return False + except InvalidToken as e: + raise ConfigEntryAuthFailed from e + except SubscriptionRequired as e: + raise ConfigEntryAuthFailed from e except TeslaFleetError as e: raise ConfigEntryNotReady from e diff --git a/homeassistant/components/teslemetry/config_flow.py b/homeassistant/components/teslemetry/config_flow.py index 0803688b1ca..5fb6ce56aed 100644 --- a/homeassistant/components/teslemetry/config_flow.py +++ b/homeassistant/components/teslemetry/config_flow.py @@ -14,7 +14,7 @@ from tesla_fleet_api.exceptions import ( ) import voluptuous as vol -from homeassistant.config_entries import ConfigFlow, ConfigFlowResult +from homeassistant.config_entries import ConfigEntry, ConfigFlow, ConfigFlowResult from homeassistant.const import CONF_ACCESS_TOKEN from homeassistant.helpers.aiohttp_client import async_get_clientsession @@ -31,33 +31,38 @@ class TeslemetryConfigFlow(ConfigFlow, domain=DOMAIN): """Config Teslemetry API connection.""" VERSION = 1 + _entry: ConfigEntry | None = None + + async def async_auth(self, user_input: Mapping[str, Any]) -> dict[str, str]: + """Reusable Auth Helper.""" + teslemetry = Teslemetry( + session=async_get_clientsession(self.hass), + access_token=user_input[CONF_ACCESS_TOKEN], + ) + try: + await teslemetry.test() + except InvalidToken: + return {CONF_ACCESS_TOKEN: "invalid_access_token"} + except SubscriptionRequired: + return {"base": "subscription_required"} + except ClientConnectionError: + return {"base": "cannot_connect"} + except TeslaFleetError as e: + LOGGER.error(e) + return {"base": "unknown"} + return {} async def async_step_user( self, user_input: Mapping[str, Any] | None = None ) -> ConfigFlowResult: """Get configuration from the user.""" errors: dict[str, str] = {} - if user_input: - teslemetry = Teslemetry( - session=async_get_clientsession(self.hass), - access_token=user_input[CONF_ACCESS_TOKEN], + if user_input and not (errors := await self.async_auth(user_input)): + self._abort_if_unique_id_configured() + return self.async_create_entry( + title="Teslemetry", + data=user_input, ) - try: - await teslemetry.test() - except InvalidToken: - errors[CONF_ACCESS_TOKEN] = "invalid_access_token" - except SubscriptionRequired: - errors["base"] = "subscription_required" - except ClientConnectionError: - errors["base"] = "cannot_connect" - except TeslaFleetError as e: - LOGGER.exception(str(e)) - errors["base"] = "unknown" - else: - return self.async_create_entry( - title="Teslemetry", - data=user_input, - ) return self.async_show_form( step_id="user", @@ -65,3 +70,31 @@ class TeslemetryConfigFlow(ConfigFlow, domain=DOMAIN): description_placeholders=DESCRIPTION_PLACEHOLDERS, errors=errors, ) + + async def async_step_reauth( + self, entry_data: Mapping[str, Any] + ) -> ConfigFlowResult: + """Handle reauth on failure.""" + 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: Mapping[str, Any] | None = None + ) -> ConfigFlowResult: + """Handle users reauth credentials.""" + + assert self._entry + errors: dict[str, str] = {} + + if user_input and not (errors := await self.async_auth(user_input)): + return self.async_update_reload_and_abort( + self._entry, + data=user_input, + ) + + return self.async_show_form( + step_id="reauth_confirm", + description_placeholders=DESCRIPTION_PLACEHOLDERS, + data_schema=TESLEMETRY_SCHEMA, + errors=errors, + ) diff --git a/homeassistant/components/teslemetry/coordinator.py b/homeassistant/components/teslemetry/coordinator.py index 27ff45f75a3..75794c7cdec 100644 --- a/homeassistant/components/teslemetry/coordinator.py +++ b/homeassistant/components/teslemetry/coordinator.py @@ -5,10 +5,15 @@ from typing import Any from tesla_fleet_api import EnergySpecific, VehicleSpecific from tesla_fleet_api.const import VehicleDataEndpoint -from tesla_fleet_api.exceptions import TeslaFleetError, VehicleOffline +from tesla_fleet_api.exceptions import ( + InvalidToken, + SubscriptionRequired, + TeslaFleetError, + VehicleOffline, +) from homeassistant.core import HomeAssistant -from homeassistant.exceptions import ConfigEntryNotReady +from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed from .const import LOGGER, TeslemetryState @@ -54,6 +59,10 @@ class TeslemetryVehicleDataCoordinator(TeslemetryDataCoordinator): if response["response"]["state"] != TeslemetryState.ONLINE: # The first refresh will fail, so retry later raise ConfigEntryNotReady("Vehicle is not online") + except InvalidToken as e: + raise ConfigEntryAuthFailed from e + except SubscriptionRequired as e: + raise ConfigEntryAuthFailed from e except TeslaFleetError as e: # The first refresh will also fail, so retry later raise ConfigEntryNotReady from e @@ -67,6 +76,10 @@ class TeslemetryVehicleDataCoordinator(TeslemetryDataCoordinator): except VehicleOffline: self.data["state"] = TeslemetryState.OFFLINE return self.data + except InvalidToken as e: + raise ConfigEntryAuthFailed from e + except SubscriptionRequired as e: + raise ConfigEntryAuthFailed from e except TeslaFleetError as e: raise UpdateFailed(e.message) from e @@ -97,6 +110,10 @@ class TeslemetryEnergyDataCoordinator(TeslemetryDataCoordinator): try: data = await self.api.live_status() + except InvalidToken as e: + raise ConfigEntryAuthFailed from e + except SubscriptionRequired as e: + raise ConfigEntryAuthFailed from e except TeslaFleetError as e: raise UpdateFailed(e.message) from e diff --git a/tests/components/teslemetry/test_config_flow.py b/tests/components/teslemetry/test_config_flow.py index f2894c695fa..2f12b202712 100644 --- a/tests/components/teslemetry/test_config_flow.py +++ b/tests/components/teslemetry/test_config_flow.py @@ -18,6 +18,10 @@ from homeassistant.data_entry_flow import FlowResultType from .const import CONFIG +from tests.common import MockConfigEntry + +BAD_CONFIG = {CONF_ACCESS_TOKEN: "bad_access_token"} + @pytest.fixture(autouse=True) def mock_test(): @@ -86,3 +90,95 @@ async def test_form_errors(hass: HomeAssistant, side_effect, error, mock_test) - CONFIG, ) assert result3["type"] is FlowResultType.CREATE_ENTRY + + +async def test_reauth(hass: HomeAssistant, mock_test) -> None: + """Test reauth flow.""" + + mock_entry = MockConfigEntry( + domain=DOMAIN, + data=BAD_CONFIG, + ) + mock_entry.add_to_hass(hass) + + result1 = await hass.config_entries.flow.async_init( + DOMAIN, + context={ + "source": config_entries.SOURCE_REAUTH, + "entry_id": mock_entry.entry_id, + }, + data=BAD_CONFIG, + ) + + assert result1["type"] is FlowResultType.FORM + assert result1["step_id"] == "reauth_confirm" + assert not result1["errors"] + + with patch( + "homeassistant.components.teslemetry.async_setup_entry", + return_value=True, + ) as mock_setup_entry: + result2 = await hass.config_entries.flow.async_configure( + result1["flow_id"], + CONFIG, + ) + await hass.async_block_till_done() + assert len(mock_setup_entry.mock_calls) == 1 + assert len(mock_test.mock_calls) == 1 + + assert result2["type"] is FlowResultType.ABORT + assert result2["reason"] == "reauth_successful" + assert mock_entry.data == CONFIG + + +@pytest.mark.parametrize( + ("side_effect", "error"), + [ + (InvalidToken, {CONF_ACCESS_TOKEN: "invalid_access_token"}), + (SubscriptionRequired, {"base": "subscription_required"}), + (ClientConnectionError, {"base": "cannot_connect"}), + (TeslaFleetError, {"base": "unknown"}), + ], +) +async def test_reauth_errors( + hass: HomeAssistant, mock_test, side_effect, error +) -> None: + """Test reauth flows that fail.""" + + # Start the reauth + mock_entry = MockConfigEntry( + domain=DOMAIN, + data=BAD_CONFIG, + ) + mock_entry.add_to_hass(hass) + + result = await hass.config_entries.flow.async_init( + DOMAIN, + context={ + "source": config_entries.SOURCE_REAUTH, + "unique_id": mock_entry.unique_id, + "entry_id": mock_entry.entry_id, + }, + data=BAD_CONFIG, + ) + + mock_test.side_effect = side_effect + result2 = await hass.config_entries.flow.async_configure( + result["flow_id"], + BAD_CONFIG, + ) + await hass.async_block_till_done() + + assert result2["type"] is FlowResultType.FORM + assert result2["errors"] == error + + # Complete the flow + mock_test.side_effect = None + result3 = await hass.config_entries.flow.async_configure( + result2["flow_id"], + CONFIG, + ) + assert "errors" not in result3 + assert result3["type"] is FlowResultType.ABORT + assert result3["reason"] == "reauth_successful" + assert mock_entry.data == CONFIG diff --git a/tests/components/teslemetry/test_init.py b/tests/components/teslemetry/test_init.py index 9742338f27a..fb405e2ee03 100644 --- a/tests/components/teslemetry/test_init.py +++ b/tests/components/teslemetry/test_init.py @@ -3,6 +3,7 @@ from datetime import timedelta from freezegun.api import FrozenDateTimeFactory +import pytest from tesla_fleet_api.exceptions import ( InvalidToken, SubscriptionRequired, @@ -20,6 +21,12 @@ from .const import WAKE_UP_ASLEEP, WAKE_UP_ONLINE from tests.common import async_fire_time_changed +ERRORS = [ + (InvalidToken, ConfigEntryState.SETUP_ERROR), + (SubscriptionRequired, ConfigEntryState.SETUP_ERROR), + (TeslaFleetError, ConfigEntryState.SETUP_RETRY), +] + async def test_load_unload(hass: HomeAssistant) -> None: """Test load and unload.""" @@ -31,28 +38,15 @@ async def test_load_unload(hass: HomeAssistant) -> None: assert entry.state is ConfigEntryState.NOT_LOADED -async def test_auth_failure(hass: HomeAssistant, mock_products) -> None: - """Test init with an authentication error.""" +@pytest.mark.parametrize(("side_effect", "state"), ERRORS) +async def test_init_error( + hass: HomeAssistant, mock_products, side_effect, state +) -> None: + """Test init with errors.""" - mock_products.side_effect = InvalidToken + mock_products.side_effect = side_effect entry = await setup_platform(hass) - assert entry.state is ConfigEntryState.SETUP_ERROR - - -async def test_subscription_failure(hass: HomeAssistant, mock_products) -> None: - """Test init with an client response error.""" - - mock_products.side_effect = SubscriptionRequired - entry = await setup_platform(hass) - assert entry.state is ConfigEntryState.SETUP_ERROR - - -async def test_other_failure(hass: HomeAssistant, mock_products) -> None: - """Test init with an client response error.""" - - mock_products.side_effect = TeslaFleetError - entry = await setup_platform(hass) - assert entry.state is ConfigEntryState.SETUP_RETRY + assert entry.state is state # Vehicle Coordinator @@ -88,11 +82,14 @@ async def test_vehicle_first_refresh( mock_vehicle_data.assert_called_once() -async def test_vehicle_first_refresh_error(hass: HomeAssistant, mock_wake_up) -> None: +@pytest.mark.parametrize(("side_effect", "state"), ERRORS) +async def test_vehicle_first_refresh_error( + hass: HomeAssistant, mock_wake_up, side_effect, state +) -> None: """Test first coordinator refresh with an error.""" - mock_wake_up.side_effect = TeslaFleetError + mock_wake_up.side_effect = side_effect entry = await setup_platform(hass) - assert entry.state is ConfigEntryState.SETUP_RETRY + assert entry.state is state async def test_vehicle_refresh_offline( @@ -111,18 +108,24 @@ async def test_vehicle_refresh_offline( mock_vehicle_data.assert_called_once() -async def test_vehicle_refresh_error(hass: HomeAssistant, mock_vehicle_data) -> None: +@pytest.mark.parametrize(("side_effect", "state"), ERRORS) +async def test_vehicle_refresh_error( + hass: HomeAssistant, mock_vehicle_data, side_effect, state +) -> None: """Test coordinator refresh with an error.""" - mock_vehicle_data.side_effect = TeslaFleetError + mock_vehicle_data.side_effect = side_effect entry = await setup_platform(hass) - assert entry.state is ConfigEntryState.SETUP_RETRY + assert entry.state is state # Test Energy Coordinator -async def test_energy_refresh_error(hass: HomeAssistant, mock_live_status) -> None: +@pytest.mark.parametrize(("side_effect", "state"), ERRORS) +async def test_energy_refresh_error( + hass: HomeAssistant, mock_live_status, side_effect, state +) -> None: """Test coordinator refresh with an error.""" - mock_live_status.side_effect = TeslaFleetError + mock_live_status.side_effect = side_effect entry = await setup_platform(hass) - assert entry.state is ConfigEntryState.SETUP_RETRY + assert entry.state is state