Add the ability to select region for Roborock (#160898)

homeassistant/components/roborock/config_flow.py

@@ -29,17 +29,24 @@ from homeassistant.const import CONF_USERNAME
 from homeassistant.core import callback
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.selector import (
+    SelectSelector,
+    SelectSelectorConfig,
+    SelectSelectorMode,
+)
 from homeassistant.helpers.service_info.dhcp import DhcpServiceInfo
 
 from . import RoborockConfigEntry
 from .const import (
     CONF_BASE_URL,
     CONF_ENTRY_CODE,
+    CONF_REGION,
     CONF_SHOW_BACKGROUND,
     CONF_USER_DATA,
     DEFAULT_DRAWABLES,
     DOMAIN,
     DRAWABLES,
+    REGION_OPTIONS,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -64,17 +71,35 @@ class RoborockFlowHandler(ConfigFlow, domain=DOMAIN):
 
         if user_input is not None:
             username = user_input[CONF_USERNAME]
+            region = user_input[CONF_REGION]
             self._username = username
             _LOGGER.debug("Requesting code for Roborock account")
+            base_url = None
+            if region != "auto":
+                base_url = f"https://{region}iot.roborock.com"
             self._client = RoborockApiClient(
-                username, session=async_get_clientsession(self.hass)
+                username,
+                base_url=base_url,
+                session=async_get_clientsession(self.hass),
             )
             errors = await self._request_code()
             if not errors:
                 return await self.async_step_code()
+
         return self.async_show_form(
             step_id="user",
-            data_schema=vol.Schema({vol.Required(CONF_USERNAME): str}),
+            data_schema=vol.Schema(
+                {
+                    vol.Required(CONF_USERNAME): str,
+                    vol.Required(CONF_REGION, default="auto"): SelectSelector(
+                        SelectSelectorConfig(
+                            options=REGION_OPTIONS,
+                            mode=SelectSelectorMode.DROPDOWN,
+                            translation_key="region",
+                        )
+                    ),
+                }
+            ),
             errors=errors,
         )
 
@@ -114,6 +139,8 @@ class RoborockFlowHandler(ConfigFlow, domain=DOMAIN):
                 user_data = await self._client.code_login_v4(code)
             except RoborockInvalidCode:
                 errors["base"] = "invalid_code"
+            except RoborockAccountDoesNotExist:
+                errors["base"] = "invalid_email_or_region"
             except RoborockException:
                 _LOGGER.exception("Unexpected exception")
                 errors["base"] = "unknown_roborock"

homeassistant/components/roborock/const.py

@@ -11,7 +11,8 @@ CONF_ENTRY_CODE = "code"
 CONF_BASE_URL = "base_url"
 CONF_USER_DATA = "user_data"
 CONF_SHOW_BACKGROUND = "show_background"
-
+CONF_REGION = "region"
+REGION_OPTIONS = ["auto", "us", "eu", "ru", "cn"]
 # Option Flow steps
 DRAWABLES = "drawables"
 

homeassistant/components/roborock/strings.json

@@ -9,6 +9,7 @@
       "invalid_code": "The code you entered was incorrect, please check it and try again.",
       "invalid_email": "There is no account associated with the email you entered, please try again.",
       "invalid_email_format": "There is an issue with the formatting of your email - please try again.",
+      "invalid_email_or_region": "Either there is no account associated with the email you entered, or there is no account in the selected region.",
       "too_frequent_code_requests": "You have attempted to request too many codes. Try again later.",
       "unknown": "[%key:common::config_flow::error::unknown%]",
       "unknown_roborock": "There was an unknown Roborock exception - please check your logs.",
@@ -30,9 +31,11 @@
       },
       "user": {
         "data": {
+          "region": "Roborock server region",
           "username": "[%key:common::config_flow::data::email%]"
         },
         "data_description": {
+          "region": "The server region your Roborock account is registered in when setting up the app. Auto is recommended unless you are having issues.",
           "username": "The email address used to sign in to the Roborock app."
         },
         "description": "Enter your Roborock email address."
@@ -545,6 +548,17 @@
       }
     }
   },
+  "selector": {
+    "region": {
+      "options": {
+        "auto": "Auto",
+        "cn": "CN",
+        "eu": "EU",
+        "ru": "RU",
+        "us": "US"
+      }
+    }
+  },
   "services": {
     "get_maps": {
       "description": "Retrieves the map and room information of your device.",

tests/components/roborock/test_config_flow.py

@@ -1,7 +1,8 @@
 """Test Roborock config flow."""
 
+import asyncio
 from copy import deepcopy
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
 import pytest
 from roborock import RoborockTooFrequentCodeRequests
@@ -15,10 +16,17 @@ from roborock.exceptions import (
 from vacuum_map_parser_base.config.drawable import Drawable
 
 from homeassistant import config_entries
-from homeassistant.components.roborock.const import CONF_ENTRY_CODE, DOMAIN, DRAWABLES
+from homeassistant.components.roborock.const import (
+    CONF_BASE_URL,
+    CONF_ENTRY_CODE,
+    CONF_REGION,
+    DOMAIN,
+    DRAWABLES,
+)
 from homeassistant.const import CONF_USERNAME, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.service_info.dhcp import DhcpServiceInfo
 
 from .mock_data import MOCK_CONFIG, NETWORK_INFO, ROBOROCK_RRUID, USER_DATA, USER_EMAIL
@@ -148,6 +156,7 @@ async def test_config_flow_failures_request_code(
     [
         (RoborockException(), {"base": "unknown_roborock"}),
         (RoborockInvalidCode(), {"base": "invalid_code"}),
+        (RoborockAccountDoesNotExist(), {"base": "invalid_email_or_region"}),
         (Exception(), {"base": "unknown"}),
     ],
 )
@@ -398,3 +407,54 @@ async def test_discovery_already_setup(
 
     assert result["type"] is FlowResultType.ABORT
     assert result["reason"] == "already_configured"
+
+
+async def test_config_flow_with_region(
+    hass: HomeAssistant,
+) -> None:
+    """Handle the config flow with a specific region."""
+    with patch(
+        "homeassistant.components.roborock.async_setup_entry", return_value=True
+    ) as mock_setup:
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN, context={"source": config_entries.SOURCE_USER}
+        )
+        assert result["type"] is FlowResultType.FORM
+        assert result["step_id"] == "user"
+        with patch(
+            "homeassistant.components.roborock.config_flow.RoborockApiClient"
+        ) as mock_client_cls:
+            mock_client = mock_client_cls.return_value
+            mock_client.request_code_v4 = AsyncMock(return_value=None)
+            mock_client.code_login_v4 = AsyncMock(return_value=USER_DATA)
+
+            # base_url is awaited in config_flow, so it needs to be an awaitable
+            future_base_url = asyncio.Future()
+            future_base_url.set_result("https://usiot.roborock.com")
+            mock_client.base_url = future_base_url
+
+            result = await hass.config_entries.flow.async_configure(
+                result["flow_id"], {CONF_USERNAME: USER_EMAIL, CONF_REGION: "us"}
+            )
+
+            # Check that the client was initialized with the correct base_url
+            mock_client_cls.assert_called_with(
+                USER_EMAIL,
+                base_url="https://usiot.roborock.com",
+                session=async_get_clientsession(hass),
+            )
+
+            assert result["type"] is FlowResultType.FORM
+            assert result["step_id"] == "code"
+            assert result["errors"] == {}
+
+            result = await hass.config_entries.flow.async_configure(
+                result["flow_id"], user_input={CONF_ENTRY_CODE: "123456"}
+            )
+
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+    assert result["context"]["unique_id"] == ROBOROCK_RRUID
+    assert result["title"] == USER_EMAIL
+    assert result["data"][CONF_BASE_URL] == "https://usiot.roborock.com"
+    assert result["result"]
+    assert len(mock_setup.mock_calls) == 1