chore: convert to store all iot login info together

Author: Lash-L

roborock/web_api.py

@@ -8,6 +8,7 @@
 import secrets
 import string
 import time
+from dataclasses import dataclass
 
 import aiohttp
 from aiohttp import ContentTypeError, FormData
@@ -31,6 +32,15 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+@dataclass
+class IotLoginInfo:
+    """Information about the login to the iot server."""
+
+    base_url: str
+    country_code: str
+    country: str
+
+
 class RoborockApiClient:
     _LOGIN_RATES = [
         Rate(1, Duration.SECOND),
@@ -48,23 +58,29 @@ class RoborockApiClient:
     _login_limiter = Limiter(_LOGIN_RATES)
     _home_data_limiter = Limiter(_HOME_DATA_RATES)
 
-    def __init__(self, username: str, base_url=None, session: aiohttp.ClientSession | None = None) -> None:
+    def __init__(
+        self, username: str, base_url: str | None = None, session: aiohttp.ClientSession | None = None
+    ) -> None:
         """Sample API Client."""
         self._username = username
-        self.base_url = base_url
+        self._base_url = base_url
         self._device_identifier = secrets.token_urlsafe(16)
         self.session = session
-        self._country = None
-        self._country_code = None
-
-    async def _get_base_url(self) -> str:
-        if not self.base_url:
-            for iot_url in [
-                "https://usiot.roborock.com",
-                "https://euiot.roborock.com",
-                "https://cniot.roborock.com",
-                "https://ruiot.roborock.com",
-            ]:
+        self._iot_login_info: IotLoginInfo | None = None
+
+    async def _get_iot_login_info(self) -> IotLoginInfo:
+        if self._iot_login_info is None:
+            valid_urls = (
+                [
+                    "https://usiot.roborock.com",
+                    "https://euiot.roborock.com",
+                    "https://cniot.roborock.com",
+                    "https://ruiot.roborock.com",
+                ]
+                if self._base_url is None
+                else [self._base_url]
+            )
+            for iot_url in valid_urls:
                 url_request = PreparedRequest(iot_url, self.session)
                 response = await url_request.request(
                     "post",
@@ -86,15 +102,31 @@ async def _get_base_url(self) -> str:
                             "Failed to get base url for %s with the following context: %s", self._username, response
                         )
                 if response["data"]["countrycode"] is not None:
-                    self._country_code = response["data"]["countrycode"]
-                    self._country = response["data"]["country"]
-                    self.base_url = response["data"]["url"]
-                    return self.base_url
+                    self._iot_login_info = IotLoginInfo(
+                        base_url=response["data"]["url"],
+                        country=response["data"]["country"],
+                        country_code=response["data"]["countrycode"],
+                    )
+                    return self._iot_login_info
             raise RoborockNoResponseFromBaseURL(
                 "No account was found for any base url we tried. Either your email is incorrect or we do not have a"
                 " record of the roborock server your device is on."
             )
-        return self.base_url
+        return self._iot_login_info
+
+    @property
+    async def base_url(self):
+        if self._base_url is not None:
+            return self._base_url
+        return (await self._get_iot_login_info()).base_url
+
+    @property
+    async def country(self):
+        return (await self._get_iot_login_info()).country
+
+    @property
+    async def country_code(self):
+        return (await self._get_iot_login_info()).country_code
 
     def _get_header_client_id(self):
         md5 = hashlib.md5()
@@ -178,7 +210,7 @@ async def request_code(self) -> None:
         except BucketFullException as ex:
             _LOGGER.info(ex.meta_info)
             raise RoborockRateLimit("Reached maximum requests for login. Please try again later.") from ex
-        base_url = await self._get_base_url()
+        base_url = await self.base_url
         header_clientid = self._get_header_client_id()
         code_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
 
@@ -209,7 +241,7 @@ async def request_code_v4(self) -> None:
         except BucketFullException as ex:
             _LOGGER.info(ex.meta_info)
             raise RoborockRateLimit("Reached maximum requests for login. Please try again later.") from ex
-        base_url = await self._get_base_url()
+        base_url = await self.base_url
         header_clientid = self._get_header_client_id()
         code_request = PreparedRequest(
             base_url,
@@ -240,7 +272,7 @@ async def request_code_v4(self) -> None:
 
     async def _sign_key_v3(self, s: str) -> str:
         """Sign a randomly generated string."""
-        base_url = await self._get_base_url()
+        base_url = await self.base_url
         header_clientid = self._get_header_client_id()
         code_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
 
@@ -269,11 +301,11 @@ async def code_login_v4(
         :param country: The two-character representation of the country, i.e. "US"
         :param country_code: the country phone number code i.e. 1 for US.
         """
-        base_url = await self._get_base_url()
+        base_url = await self.base_url
         if country is None:
-            country = self._country
+            country = await self.country
         if country_code is None:
-            country_code = self._country_code
+            country_code = await self.country_code
         header_clientid = self._get_header_client_id()
         x_mercy_ks = "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(16))
         x_mercy_k = await self._sign_key_v3(x_mercy_ks)
@@ -321,7 +353,7 @@ async def pass_login(self, password: str) -> UserData:
         except BucketFullException as ex:
             _LOGGER.info(ex.meta_info)
             raise RoborockRateLimit("Reached maximum requests for login. Please try again later.") from ex
-        base_url = await self._get_base_url()
+        base_url = await self.base_url
         header_clientid = self._get_header_client_id()
 
         login_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
@@ -360,7 +392,7 @@ async def pass_login_v3(self, password: str) -> UserData:
         raise NotImplementedError("Pass_login_v3 has not yet been implemented")
 
     async def code_login(self, code: int | str) -> UserData:
-        base_url = await self._get_base_url()
+        base_url = await self.base_url
         header_clientid = self._get_header_client_id()
 
         login_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
@@ -393,7 +425,7 @@ async def code_login(self, code: int | str) -> UserData:
         return UserData.from_dict(user_data)
 
     async def _get_home_id(self, user_data: UserData):
-        base_url = await self._get_base_url()
+        base_url = await self.base_url
         header_clientid = self._get_header_client_id()
         home_id_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
         home_id_response = await home_id_request.request(
@@ -564,7 +596,7 @@ async def execute_scene(self, user_data: UserData, scene_id: int) -> None:
 
     async def get_products(self, user_data: UserData) -> ProductResponse:
         """Gets all products and their schemas, good for determining status codes and model numbers."""
-        base_url = await self._get_base_url()
+        base_url = await self.base_url
         header_clientid = self._get_header_client_id()
         product_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
         product_response = await product_request.request(
@@ -582,7 +614,7 @@ async def get_products(self, user_data: UserData) -> ProductResponse:
         raise RoborockException("product result was an unexpected type")
 
     async def download_code(self, user_data: UserData, product_id: int):
-        base_url = await self._get_base_url()
+        base_url = await self.base_url
         header_clientid = self._get_header_client_id()
         product_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
         request = {"apilevel": 99999, "productids": [product_id], "type": 2}
@@ -595,7 +627,7 @@ async def download_code(self, user_data: UserData, product_id: int):
         return response["data"][0]["url"]
 
     async def download_category_code(self, user_data: UserData):
-        base_url = await self._get_base_url()
+        base_url = await self.base_url
         header_clientid = self._get_header_client_id()
         product_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
         response = await product_request.request(