Add region selection to Tesla Fleet config flow (#177701)

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:

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"

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"
+      }
+    }
   }
 }

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")