feat: add method for returning country and country code

roborock/web_api.py

@@ -57,32 +57,41 @@ def __init__(self, username: str, base_url=None, session: aiohttp.ClientSession
         self._device_identifier = secrets.token_urlsafe(16)
         self.session = session
 
+    async def _get_info_by_email(self, url: str) -> dict:
+        url_request = PreparedRequest(url, self.session)
+        response = await url_request.request(
+            "post",
+            "/api/v1/getUrlByEmail",
+            params={"email": self._username, "needtwostepauth": "false"},
+        )
+        if response is None:
+            raise RoborockUrlException("Get url by email returned None")
+        response_code = response.get("code")
+        if response_code != 200:
+            _LOGGER.info("Get base url failed for %s with the following context: %s", self._username, response)
+            if response_code == 2003:
+                raise RoborockInvalidEmail("Your email was incorrectly formatted.")
+            elif response_code == 1001:
+                raise RoborockMissingParameters(
+                    "You are missing parameters for this request, are you sure you entered your username?"
+                )
+            elif response_code == 9002:
+                raise RoborockTooManyRequest("Please temporarily disable making requests and try again later.")
+            raise RoborockUrlException(f"error code: {response_code} msg: {response.get('error')}")
+        response_data = response.get("data")
+        if response_data is None:
+            raise RoborockUrlException("response does not have 'data'")
+
+        return {
+            "base_url": response_data.get("url"),
+            "country": response_data.get("country"),
+            "countrycode": response_data.get("countrycode"),
+        }
+
     async def _get_base_url(self) -> str:
         if not self.base_url:
-            url_request = PreparedRequest(self._default_url, self.session)
-            response = await url_request.request(
-                "post",
-                "/api/v1/getUrlByEmail",
-                params={"email": self._username, "needtwostepauth": "false"},
-            )
-            if response is None:
-                raise RoborockUrlException("get url by email returned None")
-            response_code = response.get("code")
-            if response_code != 200:
-                _LOGGER.info("Get base url failed for %s with the following context: %s", self._username, response)
-                if response_code == 2003:
-                    raise RoborockInvalidEmail("Your email was incorrectly formatted.")
-                elif response_code == 1001:
-                    raise RoborockMissingParameters(
-                        "You are missing parameters for this request, are you sure you entered your username?"
-                    )
-                elif response_code == 9002:
-                    raise RoborockTooManyRequest("Please temporarily disable making requests and try again later.")
-                raise RoborockUrlException(f"error code: {response_code} msg: {response.get('error')}")
-            response_data = response.get("data")
-            if response_data is None:
-                raise RoborockUrlException("response does not have 'data'")
-            self.base_url = response_data.get("url")
+            data = await self._get_info_by_email(self._default_url)
+            self.base_url = data["base_url"]
         return self.base_url
 
     def _get_header_client_id(self):
@@ -161,6 +170,14 @@ async def add_device(self, user_data: UserData, s: str, t: str) -> dict:
 
         return add_device_response["result"]
 
+    async def get_country_code_and_country(self) -> dict[str, str]:
+        """Get country and country code for user."""
+
+        base_url = await self._get_base_url()
+        data = await self._get_info_by_email(base_url)
+
+        return {"country": data["country"], "countrycode": data["countrycode"]}
+
     async def request_code(self) -> None:
         try:
             self._login_limiter.try_acquire("login")

tests/conftest.py

@@ -16,7 +16,7 @@
 from roborock.roborock_message import RoborockMessage
 from roborock.version_1_apis.roborock_local_client_v1 import RoborockLocalClientV1
 from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1
-from tests.mock_data import HOME_DATA_RAW, HOME_DATA_SCENES_RAW, TEST_LOCAL_API_HOST, USER_DATA
+from tests.mock_data import BASE_URL, HOME_DATA_RAW, HOME_DATA_SCENES_RAW, TEST_LOCAL_API_HOST, USER_DATA
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -177,7 +177,7 @@ def mock_rest() -> aioresponses:
             status=200,
             payload={
                 "code": 200,
-                "data": {"country": "US", "countrycode": "1", "url": "https://usiot.roborock.com"},
+                "data": {"country": USER_DATA["country"], "countrycode": USER_DATA["countrycode"], "url": BASE_URL},
                 "msg": "success",
             },
         )

tests/test_web_api.py

@@ -71,3 +71,10 @@ async def test_code_login_v4_flow(mock_rest) -> None:
     await api.request_code_v4()
     ud = await api.code_login_v4(4123, "US", 1)
     assert ud == UserData.from_dict(USER_DATA)
+
+
+async def test_get_country_code_and_country(mock_rest) -> None:
+    """Test that we can login with a code and we get back the correct userdata object."""
+    api = RoborockApiClient(username="test_user@gmail.com", base_url="https://euiot.roborock.com")
+    res = await api.get_country_code_and_country()
+    assert res == {"country": USER_DATA["country"], "countrycode": USER_DATA["countrycode"]}