fix: fallback to old version of login (#537)

* fix: fallback to old version of login

* fix: bug and add test

Author: Lash-L

roborock/web_api.py

@@ -98,7 +98,7 @@ async def _get_iot_login_info(self) -> IotLoginInfo:
                         raise RoborockException(f"{response.get('msg')} - response code: {response_code}")
                 country_code = response["data"]["countrycode"]
                 country = response["data"]["country"]
-                if country_code is not None and country is not None:
+                if country_code is not None or country is not None:
                     self._iot_login_info = IotLoginInfo(
                         base_url=response["data"]["url"],
                         country=country,
@@ -234,6 +234,9 @@ async def request_code(self) -> None:
 
     async def request_code_v4(self) -> None:
         """Request a code using the v4 endpoint."""
+        if await self.country_code is None or await self.country is None:
+            _LOGGER.info("No country code or country found, trying old version of request code.")
+            return await self.request_code()
         try:
             self._login_limiter.try_acquire("login")
         except BucketFullException as ex:
@@ -304,6 +307,9 @@ async def code_login_v4(
             country = await self.country
         if country_code is None:
             country_code = await self.country_code
+        if country_code is None or country is None:
+            _LOGGER.info("No country code or country found, trying old version of code login.")
+            return await self.code_login(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)

tests/test_web_api.py

@@ -170,3 +170,36 @@ 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_missing_country_login(mock_rest) -> None:
+    """Test that we cycle through the URLs correctly."""
+    mock_rest.clear()
+    # Make country None, but country code set.
+    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": 1},
+            "msg": "Success",
+        },
+    )
+    # v4 is not mocked, so it would fail it were called.
+    mock_rest.post(
+        re.compile(r"https://.*iot\.roborock\.com/api/v1/loginWithCode.*"),
+        status=200,
+        payload={"code": 200, "data": USER_DATA, "msg": "success"},
+    )
+    mock_rest.post(
+        re.compile(r"https://.*iot\.roborock\.com/api/v1/sendEmailCode.*"),
+        status=200,
+        payload={"code": 200, "data": None, "msg": "success"},
+    )
+
+    client = RoborockApiClient("test@example.com")
+    await client.request_code_v4()
+    ud = await client.code_login_v4(4123)
+    assert ud is not None
+    # Ensure we have no surprise REST calls.
+    assert len(mock_rest.requests) == 3