Iterate down possible URLs when they return 3030 errors

Author: Ongy

roborock/web_api.py

@@ -9,6 +9,7 @@
 import string
 import time
 from dataclasses import dataclass
+from typing import List
 
 import aiohttp
 from aiohttp import ContentTypeError, FormData
@@ -74,11 +75,11 @@ def __init__(
         self._device_identifier = secrets.token_urlsafe(16)
         self.session = session
         self._iot_login_info: IotLoginInfo | None = None
+        self._base_urls: List[str] = BASE_URLS if base_url is None else [base_url]
 
     async def _get_iot_login_info(self) -> IotLoginInfo:
         if self._iot_login_info is None:
-            valid_urls = BASE_URLS if self._base_url is None else [self._base_url]
-            for iot_url in valid_urls:
+            for iot_url in self._base_urls:
                 url_request = PreparedRequest(iot_url, self.session)
                 response = await url_request.request(
                     "post",
@@ -95,9 +96,6 @@ async def _get_iot_login_info(self) -> IotLoginInfo:
                         raise RoborockMissingParameters(
                             "You are missing parameters for this request, are you sure you entered your username?"
                         )
-                    elif response_code == 3030:
-                        _LOGGER.warning("Found account in deletion in %s (code %s)", country, country_code)
-                        continue
                     else:
                         raise RoborockException(f"{response.get('msg')} - response code: {response_code}")
                 country_code = response["data"]["countrycode"]
@@ -272,6 +270,10 @@ async def request_code_v4(self) -> None:
                 raise RoborockAccountDoesNotExist("Account does not exist - check your login and try again.")
             elif response_code == 9002:
                 raise RoborockTooFrequentCodeRequests("You have attempted to request too many codes. Try again later")
+            elif response_code == 3030 and len(self._base_urls) > 1:
+                    self._base_urls = self._base_urls[1:]
+                    self._iot_login_info = None
+                    return await self.request_code_v4()
             else:
                 raise RoborockException(f"{code_response.get('msg')} - response code: {code_response.get('code')}")
 

tests/test_web_api.py

@@ -87,18 +87,17 @@ async def test_code_login_v4_flow(mock_rest) -> None:
     ud = await api.code_login_v4(4123, "US", 1)
     assert ud == UserData.from_dict(USER_DATA)
 
-
 async def test_url_cycling(mock_rest) -> None:
     """Test that we cycle through the URLs correctly."""
     # Clear mock rest so that we can override the patches.
     mock_rest.clear()
 
-    # 1. Mock US URL to return the account in deletion code
+    # 2. Mock US URL to return valid status but None for countrycode
     mock_rest.post(
         re.compile("https://usiot.roborock.com/api/v1/getUrlByEmail.*"),
         status=200,
         payload={
-            "code": 3030,
+            "code": 200,
             "data": {"url": "https://usiot.roborock.com", "country": None, "countrycode": None},
             "msg": "Account in deletion",
         },
@@ -183,6 +182,82 @@ async def test_url_cycling(mock_rest) -> None:
     # Make sure we just have the three we tested for above.
     assert len(mock_rest.requests) == 3
 
+async def test_thirty_thirty_cycling(mock_rest) -> None:
+    """Test that we cycle through the URLs correctly when users have deleted accounts in higher prio regions."""
+    # Clear mock rest so that we can override the patches.
+    mock_rest.clear()
+
+    mock_rest.post(
+        re.compile("https://usiot.roborock.com/api/v1/getUrlByEmail.*"),
+        status=200,
+        payload={
+            "code": 200,
+            "data": {"url": "https://usiot.roborock.com", "country": "US", "countrycode": 1},
+            "msg": "Account in deletion",
+        },
+    )
+
+    mock_rest.post(
+        re.compile("https://euiot.roborock.com/api/v1/getUrlByEmail.*"),
+        status=200,
+        payload={
+            "code": 200,
+            "data": {"url": "https://euiot.roborock.com", "country": "EU", "countrycode": 49},
+            "msg": "Success",
+        },
+    )
+
+    mock_rest.post(
+        re.compile("https://usiot.roborock.com/api/v4/email/code/send.*"),
+        status=200,
+        payload={
+            "code": 3030,
+        },
+    )
+    mock_rest.post(
+        re.compile("https://euiot.roborock.com/api/v4/email/code/send.*"),
+        status=200,
+        payload={
+            "code": 200,
+        },
+    )
+
+    mock_rest.post(re.compile("https://ruiot.roborock.com/api/v1/getUrlByEmail.*"), status=500)
+    mock_rest.post(re.compile("https://cniot.roborock.com/api/v1/getUrlByEmail.*"), status=500)
+
+    client = RoborockApiClient("test@example.com")
+    await client.request_code_v4()
+
+    print(mock_rest.requests)
+    assert (
+        len(
+            mock_rest.requests[
+                (
+                    "post",
+                    normalize_url(
+                        "https://euiot.roborock.com/api/v4/email/code/send"
+                    ),
+                )
+            ]
+        )
+        == 1
+    )
+    assert (
+        len(
+            mock_rest.requests[
+                (
+                    "post",
+                    normalize_url(
+                        "https://usiot.roborock.com/api/v4/email/code/send"
+                    ),
+                )
+            ]
+        )
+        == 1
+    )
+    # Assert that we didn't try on the Russian or Chinese regions
+    assert('https://ruiot.roborock.com/api/v4/email/code/send' not in mock_rest.requests)
+    assert('https://cniot.roborock.com/api/v4/email/code/send' not in mock_rest.requests)
 
 async def test_missing_country_login(mock_rest) -> None:
     """Test that we cycle through the URLs correctly."""