diff --git a/homeassistant/components/tesla_fleet/config_flow.py b/homeassistant/components/tesla_fleet/config_flow.py index 2285cd8dca59..3e82df6f5fdb 100644 --- a/homeassistant/components/tesla_fleet/config_flow.py +++ b/homeassistant/components/tesla_fleet/config_flow.py @@ -7,7 +7,7 @@ from typing import Any, cast, override import jwt from tesla_fleet_api import TeslaFleetApi -from tesla_fleet_api.const import SERVERS, Scope +from tesla_fleet_api.const import Scope from tesla_fleet_api.exceptions import ( InvalidToken, LoginRequired, @@ -18,16 +18,19 @@ from tesla_fleet_api.exceptions import ( import voluptuous as vol from homeassistant.config_entries import SOURCE_REAUTH, ConfigFlowResult -from homeassistant.const import CONF_DOMAIN +from homeassistant.const import CONF_DOMAIN, CONF_REGION from homeassistant.helpers import config_entry_oauth2_flow from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.selector import ( QrCodeSelector, QrCodeSelectorConfig, QrErrorCorrectionLevel, + SelectSelector, + SelectSelectorConfig, + SelectSelectorMode, ) -from .const import DOMAIN, LOGGER +from .const import DOMAIN, LOGGER, REGION_SERVERS, REGIONS from .oauth import TeslaUserImplementation @@ -44,7 +47,8 @@ class OAuth2FlowHandler( self.domain: str | None = None self.data: dict[str, Any] = {} self.uid: str | None = None - self.apis: list[TeslaFleetApi] = [] + self.region: str = REGIONS[0] + self.api: TeslaFleetApi | None = None @property @override @@ -57,7 +61,7 @@ class OAuth2FlowHandler( self, data: dict[str, Any], ) -> ConfigFlowResult: - """Handle OAuth completion and proceed to domain registration.""" + """Handle OAuth completion and proceed to region selection.""" token = jwt.decode( data["token"]["access_token"], options={"verify_signature": False} ) @@ -73,18 +77,28 @@ class OAuth2FlowHandler( ) self._abort_if_unique_id_configured() - # OAuth done, setup Partner API connections for all regions - implementation = cast(TeslaUserImplementation, self.flow_impl) - session = async_get_clientsession(self.hass) - failed_regions: list[str] = [] + # Default the region to the one detected from the OAuth token + detected = token.get("ou_code", "").lower() + self.region = detected if detected in REGIONS else REGIONS[0] - for region, server_url in SERVERS.items(): - if region == "cn": - continue + return await self.async_step_region() + + async def async_step_region( + self, user_input: dict[str, Any] | None = None + ) -> ConfigFlowResult: + """Handle region selection and partner login.""" + + errors: dict[str, str] = {} + + if user_input is not None: + self.region = user_input[CONF_REGION] + + implementation = cast(TeslaUserImplementation, self.flow_impl) + session = async_get_clientsession(self.hass) api = TeslaFleetApi( session=session, access_token="", - server=server_url, + server=REGION_SERVERS[self.region], partner_scope=True, charging_scope=False, energy_scope=False, @@ -101,29 +115,32 @@ class OAuth2FlowHandler( except (InvalidToken, OAuthExpired, LoginRequired) as err: LOGGER.warning( "Partner login failed for %s due to an authentication error: %s", - server_url, + api.server, err, ) return self.async_abort(reason="oauth_error") except TeslaFleetError as err: - LOGGER.warning("Partner login failed for %s: %s", server_url, err) - failed_regions.append(server_url) - continue - self.apis.append(api) + LOGGER.warning("Partner login failed for %s: %s", api.server, err) + errors["base"] = "cannot_connect" + else: + self.api = api + return await self.async_step_domain_input() - if not self.apis: - LOGGER.warning( - "Partner login failed for all regions: %s", ", ".join(failed_regions) - ) - return self.async_abort(reason="oauth_error") - - if failed_regions: - LOGGER.warning( - "Partner login succeeded on some regions but failed on: %s", - ", ".join(failed_regions), - ) - - return await self.async_step_domain_input() + return self.async_show_form( + step_id="region", + data_schema=vol.Schema( + { + vol.Required(CONF_REGION, default=self.region): SelectSelector( + SelectSelectorConfig( + options=REGIONS, + translation_key="region", + mode=SelectSelectorMode.LIST, + ) + ) + } + ), + errors=errors, + ) async def async_step_domain_input( self, @@ -160,40 +177,28 @@ class OAuth2FlowHandler( async def async_step_domain_registration( self, user_input: dict[str, Any] | None = None ) -> ConfigFlowResult: - """Handle domain registration for all regions.""" + """Handle domain registration for the selected region.""" - assert self.apis - assert self.apis[0].private_key + assert self.api + assert self.api.private_key assert self.domain errors: dict[str, str] = {} description_placeholders = { "public_key_url": f"https://{self.domain}/.well-known/appspecific/com.tesla.3p.public-key.pem", - "pem": self.apis[0].public_pem, + "pem": self.api.public_pem, } - successful_response: dict[str, Any] | None = None - failed_regions: list[str] = [] - - for api in self.apis: - try: - register_response = await api.partner.register(self.domain) - except PreconditionFailed: - return await self.async_step_domain_input( - errors={CONF_DOMAIN: "precondition_failed"} - ) - except TeslaFleetError as e: - LOGGER.warning( - "Partner registration failed for %s: %s", - api.server, - e.message, - ) - failed_regions.append(api.server or "unknown") - else: - if successful_response is None: - successful_response = register_response - - if successful_response is None: + try: + register_response = await self.api.partner.register(self.domain) + except PreconditionFailed: + return await self.async_step_domain_input( + errors={CONF_DOMAIN: "precondition_failed"} + ) + except TeslaFleetError as e: + LOGGER.warning( + "Partner registration failed for %s: %s", self.api.server, e.message + ) errors["base"] = "invalid_response" return self.async_show_form( step_id="domain_registration", @@ -201,22 +206,12 @@ class OAuth2FlowHandler( errors=errors, ) - if failed_regions: - LOGGER.warning( - "Partner registration succeeded on some regions but failed on: %s", - ", ".join(failed_regions), - ) - - # Verify public key from the successful response - registered_public_key = successful_response.get("response", {}).get( - "public_key" - ) + registered_public_key = register_response.get("response", {}).get("public_key") if not registered_public_key: errors["base"] = "public_key_not_found" elif ( - registered_public_key.lower() - != self.apis[0].public_uncompressed_point.lower() + registered_public_key.lower() != self.api.public_uncompressed_point.lower() ): errors["base"] = "public_key_mismatch" else: diff --git a/homeassistant/components/tesla_fleet/const.py b/homeassistant/components/tesla_fleet/const.py index ed524b634856..2670d6199594 100644 --- a/homeassistant/components/tesla_fleet/const.py +++ b/homeassistant/components/tesla_fleet/const.py @@ -3,12 +3,18 @@ from enum import StrEnum import logging -from tesla_fleet_api.const import Scope +from tesla_fleet_api.const import SERVERS, Scope DOMAIN = "tesla_fleet" CONF_REFRESH_TOKEN = "refresh_token" +# Regions the user can register in; China uses separate infrastructure. +REGION_SERVERS: dict[str, str] = { + region: server for region, server in SERVERS.items() if region != "cn" +} +REGIONS = list(REGION_SERVERS) + LOGGER = logging.getLogger(__package__) AUTHORIZE_URL = "https://fleet-auth.prd.vn.cloud.tesla.com/oauth2/v3/authorize" diff --git a/homeassistant/components/tesla_fleet/strings.json b/homeassistant/components/tesla_fleet/strings.json index 8b9cfe120569..576eca03e747 100644 --- a/homeassistant/components/tesla_fleet/strings.json +++ b/homeassistant/components/tesla_fleet/strings.json @@ -18,6 +18,7 @@ "default": "Successfully authenticated with Tesla." }, "error": { + "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]", "invalid_domain": "Invalid domain format. Please enter a valid domain name.", "invalid_response": "The registration was rejected by Tesla", "precondition_failed": "The domain does not match the application's allowed origins.", @@ -53,6 +54,15 @@ "description": "The {name} integration needs to re-authenticate your account. Reauthentication refreshes the Tesla API permissions granted to Home Assistant, including any newly enabled scopes.", "title": "[%key:common::config_flow::title::reauth%]" }, + "region": { + "data": { + "region": "Region" + }, + "data_description": { + "region": "The region to register your application in. This defaults to the region detected from your Tesla account and only needs changing if your vehicles or energy sites are in a different region." + }, + "title": "Select region" + }, "registration_complete": { "data": { "qr_code": "QR code" @@ -645,5 +655,13 @@ "wake_up_timeout": { "message": "Could not wake up vehicle" } + }, + "selector": { + "region": { + "options": { + "eu": "Europe, Middle East & Africa", + "na": "North America & Asia-Pacific" + } + } } } diff --git a/tests/components/tesla_fleet/test_config_flow.py b/tests/components/tesla_fleet/test_config_flow.py index 07829d7740dc..bcf532a50862 100644 --- a/tests/components/tesla_fleet/test_config_flow.py +++ b/tests/components/tesla_fleet/test_config_flow.py @@ -4,6 +4,7 @@ from unittest.mock import AsyncMock, Mock, patch from urllib.parse import parse_qs, urlparse import pytest +from tesla_fleet_api.const import SERVERS from tesla_fleet_api.exceptions import ( InvalidResponse, LoginRequired, @@ -24,7 +25,7 @@ from homeassistant.components.tesla_fleet.const import ( TOKEN_URL, ) from homeassistant.config_entries import SOURCE_USER -from homeassistant.const import CONF_DOMAIN +from homeassistant.const import CONF_DOMAIN, CONF_REGION from homeassistant.core import HomeAssistant from homeassistant.data_entry_flow import FlowResultType from homeassistant.helpers import config_entry_oauth2_flow @@ -137,20 +138,27 @@ async def test_partner_login_auth_error( mock_api_class.return_value = mock_api result = await hass.config_entries.flow.async_configure(result["flow_id"]) + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "region" + + # Selecting a region triggers partner login, which fails to authenticate + result = await hass.config_entries.flow.async_configure( + result["flow_id"], {CONF_REGION: "na"} + ) assert result["type"] is FlowResultType.ABORT assert result["reason"] == "oauth_error" @pytest.mark.usefixtures("current_request_with_host") -async def test_partner_login_partial_failure( +async def test_region_partner_login_error( hass: HomeAssistant, hass_client_no_auth: ClientSessionGenerator, aioclient_mock: AiohttpClientMocker, access_token: str, mock_private_key, ) -> None: - """Test partner login succeeds when one region fails.""" + """Test a partner login error keeps the user on the region step.""" result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": SOURCE_USER} ) @@ -183,34 +191,38 @@ async def test_partner_login_partial_failure( "112233445566778899aabbccddeeff1122" ) - mock_api_na = AsyncMock() - mock_api_na.private_key = mock_private_key - mock_api_na.get_private_key = AsyncMock() - mock_api_na.partner_login = AsyncMock() - mock_api_na.public_uncompressed_point = public_key - mock_api_na.partner.register.return_value = {"response": {"public_key": public_key}} - - mock_api_eu = AsyncMock() - mock_api_eu.private_key = mock_private_key - mock_api_eu.get_private_key = AsyncMock() - mock_api_eu.partner_login = AsyncMock( - side_effect=TeslaFleetError("EU partner login failed") + mock_api = AsyncMock() + mock_api.private_key = mock_private_key + mock_api.get_private_key = AsyncMock() + mock_api.partner_login = AsyncMock( + side_effect=[TeslaFleetError("Partner login failed"), None] ) + mock_api.public_uncompressed_point = public_key + mock_api.partner.register.return_value = {"response": {"public_key": public_key}} with patch( "homeassistant.components.tesla_fleet.config_flow.TeslaFleetApi", - side_effect=[mock_api_na, mock_api_eu], + return_value=mock_api, ): result = await hass.config_entries.flow.async_configure(result["flow_id"]) assert result["type"] is FlowResultType.FORM - assert result["step_id"] == "domain_input" + assert result["step_id"] == "region" + # Partner login fails, stay on the region step with an error result = await hass.config_entries.flow.async_configure( - result["flow_id"], {CONF_DOMAIN: "example.com"} + result["flow_id"], {CONF_REGION: "na"} + ) + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "region" + assert result["errors"] == {"base": "cannot_connect"} + + # Retrying succeeds and advances to the domain step + result = await hass.config_entries.flow.async_configure( + result["flow_id"], {CONF_REGION: "na"} ) assert result["type"] is FlowResultType.FORM - assert result["step_id"] == "registration_complete" + assert result["step_id"] == "domain_input" @pytest.mark.usefixtures("current_request_with_host") @@ -296,6 +308,19 @@ async def test_full_flow_with_domain_registration( # Complete OAuth result = await hass.config_entries.flow.async_configure(result["flow_id"]) assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "region" + + # The region detected from the token is the default selection + region_key = next( + key for key in result["data_schema"].schema if key == CONF_REGION + ) + assert region_key.default() == "na" + + # Accept the detected region + result = await hass.config_entries.flow.async_configure( + result["flow_id"], {CONF_REGION: "na"} + ) + assert result["type"] is FlowResultType.FORM assert result["step_id"] == "domain_input" # Enter domain - this should automatically register and go to @@ -312,6 +337,8 @@ async def test_full_flow_with_domain_registration( assert result["type"] is FlowResultType.CREATE_ENTRY assert result["title"] == UNIQUE_ID assert result["result"].unique_id == UNIQUE_ID + # The selected region determines which server is registered + assert mock_api_class.call_args.kwargs["server"] == SERVERS["na"] @pytest.mark.usefixtures("current_request_with_host") @@ -362,6 +389,13 @@ async def test_domain_input_invalid_domain( # Complete OAuth result = await hass.config_entries.flow.async_configure(result["flow_id"]) assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "region" + + # Accept the detected region + result = await hass.config_entries.flow.async_configure( + result["flow_id"], {CONF_REGION: "na"} + ) + assert result["type"] is FlowResultType.FORM assert result["step_id"] == "domain_input" # Enter invalid domain @@ -455,6 +489,13 @@ async def test_domain_registration_errors( # Complete OAuth result = await hass.config_entries.flow.async_configure(result["flow_id"]) + assert result["step_id"] == "region" + + # Accept the detected region + result = await hass.config_entries.flow.async_configure( + result["flow_id"], {CONF_REGION: "na"} + ) + assert result["step_id"] == "domain_input" # Enter domain - this should fail and stay on domain_registration result = await hass.config_entries.flow.async_configure( @@ -514,6 +555,13 @@ async def test_domain_registration_precondition_failed( # Complete OAuth result = await hass.config_entries.flow.async_configure(result["flow_id"]) + assert result["step_id"] == "region" + + # Accept the detected region + result = await hass.config_entries.flow.async_configure( + result["flow_id"], {CONF_REGION: "na"} + ) + assert result["step_id"] == "domain_input" # Enter domain - this should go to domain_registration # and then fail back to domain_input @@ -574,6 +622,13 @@ async def test_domain_registration_public_key_not_found( # Complete OAuth result = await hass.config_entries.flow.async_configure(result["flow_id"]) + assert result["step_id"] == "region" + + # Accept the detected region + result = await hass.config_entries.flow.async_configure( + result["flow_id"], {CONF_REGION: "na"} + ) + assert result["step_id"] == "domain_input" # Enter domain - this should fail and stay on domain_registration result = await hass.config_entries.flow.async_configure( @@ -635,6 +690,13 @@ async def test_domain_registration_public_key_mismatch( # Complete OAuth result = await hass.config_entries.flow.async_configure(result["flow_id"]) + assert result["step_id"] == "region" + + # Accept the detected region + result = await hass.config_entries.flow.async_configure( + result["flow_id"], {CONF_REGION: "na"} + ) + assert result["step_id"] == "domain_input" # Enter domain - this should fail and stay on domain_registration result = await hass.config_entries.flow.async_configure( @@ -646,14 +708,14 @@ async def test_domain_registration_public_key_mismatch( @pytest.mark.usefixtures("current_request_with_host") -async def test_domain_registration_partial_failure( +async def test_region_override( hass: HomeAssistant, hass_client_no_auth: ClientSessionGenerator, aioclient_mock: AiohttpClientMocker, access_token: str, mock_private_key, ) -> None: - """Test domain registration succeeds when one region fails.""" + """Test overriding the detected region registers using the selected region.""" result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": SOURCE_USER} ) @@ -686,61 +748,64 @@ async def test_domain_registration_partial_failure( "112233445566778899aabbccddeeff1122" ) - # Create two separate mocks for NA and EU - mock_api_na = AsyncMock() - mock_api_na.private_key = mock_private_key - mock_api_na.get_private_key = AsyncMock() - mock_api_na.partner_login = AsyncMock() - mock_api_na.public_pem = "test_pem" - mock_api_na.public_uncompressed_point = public_key - mock_api_na.partner.register.return_value = {"response": {"public_key": public_key}} - - mock_api_eu = AsyncMock() - mock_api_eu.private_key = mock_private_key - mock_api_eu.get_private_key = AsyncMock() - mock_api_eu.partner_login = AsyncMock() - mock_api_eu.public_pem = "test_pem" - mock_api_eu.public_uncompressed_point = public_key - mock_api_eu.server = "https://fleet-api.prd.eu.vn.cloud.tesla.com" - mock_api_eu.partner.register.side_effect = TeslaFleetError("EU registration failed") - with ( patch( - "homeassistant.components.tesla_fleet.config_flow.TeslaFleetApi", - side_effect=[mock_api_na, mock_api_eu], - ), + "homeassistant.components.tesla_fleet.config_flow.TeslaFleetApi" + ) as mock_api_class, patch( "homeassistant.components.tesla_fleet.async_setup_entry", return_value=True ), ): + mock_api = AsyncMock() + mock_api.private_key = mock_private_key + mock_api.get_private_key = AsyncMock() + mock_api.partner_login = AsyncMock() + mock_api.public_uncompressed_point = public_key + mock_api.partner.register.return_value = { + "response": {"public_key": public_key} + } + mock_api_class.return_value = mock_api + # Complete OAuth result = await hass.config_entries.flow.async_configure(result["flow_id"]) - assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "region" + + # Override the detected region (NA) with EU + result = await hass.config_entries.flow.async_configure( + result["flow_id"], {CONF_REGION: "eu"} + ) assert result["step_id"] == "domain_input" - # Enter domain - NA succeeds, EU fails, should still proceed result = await hass.config_entries.flow.async_configure( result["flow_id"], {CONF_DOMAIN: "example.com"} ) - assert result["type"] is FlowResultType.FORM assert result["step_id"] == "registration_complete" - # Complete flow result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) assert result["type"] is FlowResultType.CREATE_ENTRY assert result["title"] == UNIQUE_ID + # The overridden region determines which server is registered + assert mock_api_class.call_args.kwargs["server"] == SERVERS["eu"] @pytest.mark.usefixtures("current_request_with_host") -async def test_domain_registration_all_regions_fail( +async def test_region_default_fallback( hass: HomeAssistant, hass_client_no_auth: ClientSessionGenerator, aioclient_mock: AiohttpClientMocker, - access_token: str, - mock_private_key, ) -> None: - """Test domain registration fails when all regions fail.""" + """Test the region defaults to a selectable region when the token region is not.""" + token = config_entry_oauth2_flow._encode_jwt( + hass, + { + "sub": UNIQUE_ID, + "aud": [], + "scp": ["openid", "offline_access"], + "ou_code": "CN", + }, + ) + result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": SOURCE_USER} ) @@ -760,47 +825,18 @@ async def test_domain_registration_all_regions_fail( TOKEN_URL, json={ "refresh_token": "mock-refresh-token", - "access_token": access_token, + "access_token": token, "type": "Bearer", "expires_in": 60, }, ) - mock_api_na = AsyncMock() - mock_api_na.private_key = mock_private_key - mock_api_na.get_private_key = AsyncMock() - mock_api_na.partner_login = AsyncMock() - mock_api_na.public_pem = "test_pem" - mock_api_na.public_uncompressed_point = "test_point" - mock_api_na.server = "https://fleet-api.prd.na.vn.cloud.tesla.com" - mock_api_na.partner.register.side_effect = TeslaFleetError("NA registration failed") + result = await hass.config_entries.flow.async_configure(result["flow_id"]) + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "region" - mock_api_eu = AsyncMock() - mock_api_eu.private_key = mock_private_key - mock_api_eu.get_private_key = AsyncMock() - mock_api_eu.partner_login = AsyncMock() - mock_api_eu.public_pem = "test_pem" - mock_api_eu.public_uncompressed_point = "test_point" - mock_api_eu.server = "https://fleet-api.prd.eu.vn.cloud.tesla.com" - mock_api_eu.partner.register.side_effect = TeslaFleetError("EU registration failed") - - with patch( - "homeassistant.components.tesla_fleet.config_flow.TeslaFleetApi", - side_effect=[mock_api_na, mock_api_eu], - ): - # Complete OAuth - result = await hass.config_entries.flow.async_configure(result["flow_id"]) - - # Enter domain - both regions fail - with patch( - "homeassistant.helpers.translation.async_get_translations", return_value={} - ): - result = await hass.config_entries.flow.async_configure( - result["flow_id"], {CONF_DOMAIN: "example.com"} - ) - assert result["type"] is FlowResultType.FORM - assert result["step_id"] == "domain_registration" - assert result["errors"] == {"base": "invalid_response"} + region_key = next(key for key in result["data_schema"].schema if key == CONF_REGION) + assert region_key.default() == "na" @pytest.mark.usefixtures("current_request_with_host")