Change: Allow RemoteProtocolError to trigger retry (#1238)

* Change: Allow RemoteProtocolError to trigger retry

* Add: Unit tests

* Change: Variable name

Author: n-thumann

pontos/nvd/api.py

@@ -23,7 +23,13 @@
     TypeVar,
 )
 
-from httpx import URL, AsyncClient, Response, Timeout
+from httpx import (
+    URL,
+    AsyncClient,
+    RemoteProtocolError,
+    Response,
+    Timeout,
+)
 
 from pontos.errors import PontosError
 from pontos.helper import snake_case
@@ -450,19 +456,34 @@ async def _get(
         """
         headers = self._request_headers()
 
+        latest_error: Response | RemoteProtocolError
+
         for attempt in range(self._request_attempts):
             if attempt > 0:
                 delay = RETRY_DELAY**attempt
                 await asyncio.sleep(delay)
 
             await self._consider_rate_limit()
-            response = await self._client.get(
-                self._url, headers=headers, params=params
-            )
-            if not response.is_server_error:
-                break
 
-        return response
+            try:
+                response = await self._client.get(
+                    self._url, headers=headers, params=params
+                )
+
+                if response.is_server_error:
+                    latest_error = response
+                    continue
+
+                return response
+
+            except RemoteProtocolError as e:
+                latest_error = e
+                continue
+
+        if isinstance(latest_error, RemoteProtocolError):
+            raise latest_error
+
+        return latest_error
 
     async def __aenter__(self) -> "NVDApi":
         # reset rate limit counter

tests/nvd/test_api.py

@@ -11,7 +11,7 @@
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock, call, patch
 
-from httpx import AsyncClient, Response
+from httpx import AsyncClient, RemoteProtocolError, Response
 
 from pontos.nvd.api import (
     JSON,
@@ -206,6 +206,43 @@ async def test_no_retry(
         sleep_mock.assert_not_called()
         self.assertFalse(result.is_server_error)
 
+    @patch("pontos.nvd.api.asyncio.sleep", autospec=True)
+    @patch("pontos.nvd.api.AsyncClient", spec=AsyncClient)
+    async def test_exceed_attempts(
+        self,
+        async_client: MagicMock,
+        sleep_mock: MagicMock,
+    ):
+        response_mocks = [
+            RemoteProtocolError("RIP connection"),
+            MagicMock(spec=Response, is_server_error=True),
+        ]
+        http_client = AsyncMock()
+        http_client.get.side_effect = response_mocks
+        async_client.return_value = http_client
+
+        api = NVDApi("https://foo.bar/baz", request_attempts=2)
+
+        result = await api._get()
+
+        calls = [call(2.0)]
+        sleep_mock.assert_has_calls(calls)
+        self.assertIsInstance(result, Response)
+
+    @patch("pontos.nvd.api.AsyncClient", spec=AsyncClient)
+    async def test_remote_protocol_error(
+        self,
+        async_client: MagicMock,
+    ):
+        http_client = AsyncMock()
+        http_client.get.side_effect = RemoteProtocolError("RIP connection")
+        async_client.return_value = http_client
+
+        api = NVDApi("https://foo.bar/baz")
+
+        with self.assertRaises(RemoteProtocolError):
+            await api._get()
+
 
 class Result:
     def __init__(self, value: int) -> None: