chore: add tests

Author: Lash-L

roborock/web_api.py

@@ -30,6 +30,12 @@
 )
 
 _LOGGER = logging.getLogger(__name__)
+BASE_URLS = [
+    "https://usiot.roborock.com",
+    "https://euiot.roborock.com",
+    "https://cniot.roborock.com",
+    "https://ruiot.roborock.com",
+]
 
 
 @dataclass
@@ -70,16 +76,7 @@ def __init__(
 
     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]
-            )
+            valid_urls = BASE_URLS 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(
@@ -98,9 +95,7 @@ async def _get_iot_login_info(self) -> IotLoginInfo:
                             "You are missing parameters for this request, are you sure you entered your username?"
                         )
                     else:
-                        _LOGGER.warning(
-                            "Failed to get base url for %s with the following context: %s", self._username, response
-                        )
+                        raise RoborockException(f"{response.get('msg')} - response code: {response_code}")
                 if response["data"]["countrycode"] is not None:
                     self._iot_login_info = IotLoginInfo(
                         base_url=response["data"]["url"],

tests/test_web_api.py

@@ -1,7 +1,10 @@
+import re
+
 import aiohttp
+from aioresponses.compat import normalize_url
 
 from roborock import HomeData, HomeDataScene, UserData
-from roborock.web_api import RoborockApiClient
+from roborock.web_api import IotLoginInfo, RoborockApiClient
 from tests.mock_data import HOME_DATA_RAW, USER_DATA
 
 
@@ -71,3 +74,99 @@ 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_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 valid status but None for countrycode
+
+    mock_rest.post(
+        re.compile("https://usiot.roborock.com/api/v1/getUrlByEmail.*"),
+        status=200,
+        payload={
+            "code": 200,
+            "data": {"url": "https://usiot.roborock.com", "country": None, "countrycode": None},
+            "msg": "Success",
+        },
+    )
+
+    # 2. Mock EU URL to return a server-side error code
+    mock_rest.post(
+        re.compile("https://euiot.roborock.com/api/v1/getUrlByEmail.*"),
+        status=200,
+        payload={
+            "code": 200,
+            "data": {"url": "https://euiot.roborock.com", "country": None, "countrycode": None},
+            "msg": "Success",
+        },
+    )
+
+    # 3. Mock CN URL to return the correct, valid data
+    mock_rest.post(
+        re.compile("https://cniot.roborock.com/api/v1/getUrlByEmail.*"),
+        status=200,
+        payload={
+            "code": 200,
+            "data": {"url": "https://cniot.roborock.com", "country": "CN", "countrycode": "86"},
+            "msg": "Success",
+        },
+    )
+
+    # The RU URL should not be called, but we can mock it just in case
+    # to catch unexpected behavior.
+    mock_rest.post(re.compile("https://ruiot.roborock.com/api/v1/getUrlByEmail.*"), status=500)
+
+    client = RoborockApiClient("test@example.com")
+    result = await client._get_iot_login_info()
+
+    assert result is not None
+    assert isinstance(result, IotLoginInfo)
+    assert result.base_url == "https://cniot.roborock.com"
+    assert result.country == "CN"
+    assert result.country_code == "86"
+
+    assert client._iot_login_info == result
+    # Check that all three urls were called. We have to do this kind of weirdly as aioresponses seems to have a bug.
+    assert (
+        len(
+            mock_rest.requests[
+                (
+                    "post",
+                    normalize_url(
+                        "https://usiot.roborock.com/api/v1/getUrlByEmail?email=test%2540example.com&needtwostepauth=false"
+                    ),
+                )
+            ]
+        )
+        == 1
+    )
+    assert (
+        len(
+            mock_rest.requests[
+                (
+                    "post",
+                    normalize_url(
+                        "https://euiot.roborock.com/api/v1/getUrlByEmail?email=test%2540example.com&needtwostepauth=false"
+                    ),
+                )
+            ]
+        )
+        == 1
+    )
+    assert (
+        len(
+            mock_rest.requests[
+                (
+                    "post",
+                    normalize_url(
+                        "https://cniot.roborock.com/api/v1/getUrlByEmail?email=test%2540example.com&needtwostepauth=false"
+                    ),
+                )
+            ]
+        )
+        == 1
+    )
+    # Make sure we just have the three we tested for above.
+    assert len(mock_rest.requests) == 3