Fix handling 401 after 429 status code (#585)

* Fix handling 401 after 429 status code, increase retry times slightly

* Raise on multiple 401

* Fix mypy

bimmer_connected/api/authentication.py

@@ -85,31 +85,35 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
         # Try getting a response
         response: httpx.Response = (yield request)
         await response.aread()
+        prev_response_code: int = 0
 
-        # Handle "classic" 401 Unauthorized
-        if response.status_code == 401:
-            async with self.login_lock:
-                _LOGGER.debug("Received unauthorized response, refreshing token.")
-                await self.login()
-            request.headers["authorization"] = f"Bearer {self.access_token}"
-            request.headers["bmw-session-id"] = self.session_id
-            yield request
-        # Quota errors can either be 429 Too Many Requests or 403 Quota Exceeded (instead of 403 Forbidden)
-        elif response.status_code == 429 or (response.status_code == 403 and "quota" in response.text.lower()):
-            for _ in range(3):
-                if response.status_code == 429:
-                    await response.aread()
-                    wait_time = math.ceil(
-                        next(iter([int(i) for i in response.json().get("message", "") if i.isdigit()]), 2) * 1.25
-                    )
-                    _LOGGER.debug("Sleeping %s seconds due to 429 Too Many Requests", wait_time)
-                    await asyncio.sleep(wait_time)
-                    response = yield request
-            # Raise if still error after 3rd retry
-            try:
-                response.raise_for_status()
-            except httpx.HTTPStatusError as ex:
-                await handle_httpstatuserror(ex, log_handler=_LOGGER)
+        # Retry 3 times on 401 or 429
+        for _ in range(3):
+            # Handle "classic" 401 Unauthorized and try getting a new token
+            # We don't want to call the auth endpoint too many times, so we only do it once per 401
+            if response.status_code == 401 and response.status_code != prev_response_code:
+                prev_response_code = response.status_code
+                async with self.login_lock:
+                    _LOGGER.debug("Received unauthorized response, refreshing token.")
+                    await self.login()
+                request.headers["authorization"] = f"Bearer {self.access_token}"
+                request.headers["bmw-session-id"] = self.session_id
+                response = yield request
+
+            # Quota errors can either be 429 Too Many Requests or 403 Quota Exceeded (instead of 403 Forbidden)
+            elif response.status_code == 429 or (response.status_code == 403 and "quota" in response.text.lower()):
+                prev_response_code = response.status_code
+                await response.aread()
+                wait_time = get_retry_wait_time(response)
+                _LOGGER.debug("Sleeping %s seconds due to 429 Too Many Requests", wait_time)
+                await asyncio.sleep(wait_time)
+                response = yield request
+
+        # Raise if request still was not successful
+        try:
+            response.raise_for_status()
+        except httpx.HTTPStatusError as ex:
+            await handle_httpstatuserror(ex, log_handler=_LOGGER)
 
     async def login(self) -> None:
         """Get a valid OAuth token."""
@@ -399,9 +403,7 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
         for _ in range(3):
             if response.status_code == 429:
                 await response.aread()
-                wait_time = math.ceil(
-                    next(iter([int(i) for i in response.json().get("message", "") if i.isdigit()]), 2) * 1.25
-                )
+                wait_time = get_retry_wait_time(response)
                 _LOGGER.debug("Sleeping %s seconds due to 429 Too Many Requests", wait_time)
                 await asyncio.sleep(wait_time)
                 response = yield request
@@ -412,3 +414,13 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                 response.raise_for_status()
             except httpx.HTTPStatusError as ex:
                 await handle_httpstatuserror(ex, log_handler=_LOGGER)
+
+
+def get_retry_wait_time(response: httpx.Response) -> int:
+    """Get the wait time for the next retry from the response and multiply by 2."""
+    try:
+        response_wait_time = next(iter([int(i) for i in response.json().get("message", "") if i.isdigit()]))
+    except Exception:
+        response_wait_time = 2
+    wait_time = math.ceil(response_wait_time * 2)
+    return wait_time

bimmer_connected/tests/test_account.py

@@ -476,6 +476,63 @@ async def test_429_retry_raise_vehicles(caplog, bmw_fixture: respx.Router):
     assert len(log_429) == 3
 
 
+@pytest.mark.asyncio
+async def test_429_retry_with_login_ok_vehicles(bmw_fixture: respx.Router):
+    """Test the login flow but experiencing a 429 first."""
+    account = MyBMWAccount(TEST_USERNAME, TEST_PASSWORD, TEST_REGION)
+
+    json_429 = {"statusCode": 429, "message": "Rate limit is exceeded. Try again in 2 seconds."}
+
+    bmw_fixture.get("/eadrax-vcs/v4/vehicles").mock(
+        side_effect=[
+            httpx.Response(429, json=json_429),
+            httpx.Response(429, json=json_429),
+            *[httpx.Response(200, json=[])] * 2,
+        ]
+    )
+
+    with mock.patch("asyncio.sleep", new_callable=mock.AsyncMock):
+        await account.get_vehicles()
+
+
+@pytest.mark.asyncio
+async def test_429_retry_with_login_raise_vehicles(bmw_fixture: respx.Router):
+    """Test the error handling, experiencing a 429, 401 and another two 429."""
+    account = MyBMWAccount(TEST_USERNAME, TEST_PASSWORD, TEST_REGION)
+
+    json_429 = {"statusCode": 429, "message": "Rate limit is exceeded. Try again in 2 seconds."}
+
+    bmw_fixture.get("/eadrax-vcs/v4/vehicles").mock(
+        side_effect=[
+            httpx.Response(429, json=json_429),
+            httpx.Response(401),
+            httpx.Response(429, json=json_429),
+            httpx.Response(429, json=json_429),
+        ]
+    )
+
+    with mock.patch("asyncio.sleep", new_callable=mock.AsyncMock):
+        with pytest.raises(MyBMWQuotaError):
+            await account.get_vehicles()
+
+
+@pytest.mark.asyncio
+async def test_multiple_401(bmw_fixture: respx.Router):
+    """Test the error handling, when multiple 401 are received in sequence."""
+    account = MyBMWAccount(TEST_USERNAME, TEST_PASSWORD, TEST_REGION)
+
+    bmw_fixture.get("/eadrax-vcs/v4/vehicles").mock(
+        side_effect=[
+            httpx.Response(401),
+            httpx.Response(401),
+        ]
+    )
+
+    with mock.patch("asyncio.sleep", new_callable=mock.AsyncMock):
+        with pytest.raises(MyBMWAuthError):
+            await account.get_vehicles()
+
+
 @pytest.mark.asyncio
 async def test_403_quota_exceeded_vehicles_usa(caplog, bmw_fixture: respx.Router):
     """Test 403 quota issues for vehicle state and fail if it happens too often."""

bimmer_connected/tests/test_api.py

@@ -1,10 +1,12 @@
 """Tests for API that are not covered by other tests."""
 import json
 
+import httpx
 import pytest
 import respx
 
 from bimmer_connected.account import MyBMWAccount
+from bimmer_connected.api.authentication import get_retry_wait_time
 from bimmer_connected.api.regions import get_region_from_name, valid_regions
 from bimmer_connected.api.utils import anonymize_data
 from bimmer_connected.utils import log_response_store_to_file
@@ -120,3 +122,19 @@ async def test_fingerprint_deque(monkeypatch: pytest.MonkeyPatch, bmw_fixture: r
     account.config.set_log_responses(True)
     await account.get_vehicles()
     assert len(account.get_stored_responses()) == 10
+
+
+def test_get_retry_wait_time():
+    """Test extraction of retry wait time."""
+
+    # Parsing correctly
+    r = httpx.Response(429, json={"statusCode": 429, "message": "Rate limit is exceeded. Try again in 1 seconds."})
+    assert get_retry_wait_time(r) == 2
+
+    # No number found
+    r = httpx.Response(429, json={"statusCode": 429, "message": "Rate limit is exceeded."})
+    assert get_retry_wait_time(r) == 4
+
+    # No JSON response
+    r = httpx.Response(429, text="Rate limit is exceeded.")
+    assert get_retry_wait_time(r) == 4