Add Auth service to RetryingClientSession class

HardNorth · HardNorth · commit ab357dc89e91 · 2025-11-05T15:50:41.000+03:00
diff --git a/reportportal_client/_internal/aio/http.py b/reportportal_client/_internal/aio/http.py
@@ -29,10 +29,13 @@
 from aenum import Enum
 from aiohttp import ClientResponse, ClientResponseError, ClientSession, ServerConnectionError
 
+from reportportal_client._internal.services.auth import AuthAsync
+
 DEFAULT_RETRY_NUMBER: int = 5
 DEFAULT_RETRY_DELAY: float = 0.005
 THROTTLING_STATUSES: set = {425, 429}
 RETRY_STATUSES: set = {408, 500, 502, 503, 507}.union(THROTTLING_STATUSES)
+AUTH_PROBLEM_STATUSES: set = {401, 403}
 
 
 class RetryClass(int, Enum):
@@ -49,28 +52,32 @@ class RetryingClientSession:
     _client: ClientSession
     __retry_number: int
     __retry_delay: float
+    __auth: Optional[AuthAsync]
 
     def __init__(
         self,
         *args,
         max_retry_number: int = DEFAULT_RETRY_NUMBER,
         base_retry_delay: float = DEFAULT_RETRY_DELAY,
+        auth: Optional[AuthAsync] = None,
         **kwargs,
     ):
         """Initialize an instance of the session with arguments.
 
         To obtain the full list of arguments please see aiohttp.ClientSession.__init__() method. This class
-        just bypass everything to the base method, except two local arguments 'max_retry_number' and
-        'base_retry_delay'.
+        just bypass everything to the base method, except three local arguments 'max_retry_number',
+        'base_retry_delay', and 'auth'.
 
         :param max_retry_number: the maximum number of the request retries if it was unsuccessful
         :param base_retry_delay: base value for retry delay, determine how much time the class will wait after
                                  an error. Real value highly depends on Retry Class and Retry attempt number,
                                  since retries are performed in exponential delay manner
+        :param auth:             authentication instance to use for requests
         """
         self._client = ClientSession(*args, **kwargs)
         self.__retry_number = max_retry_number
         self.__retry_delay = base_retry_delay
+        self.__auth = auth
 
     async def __nothing(self):
         pass
@@ -89,12 +96,24 @@ async def __request(self, method: Callable, url, **kwargs: Any) -> ClientRespons
         400 Bad Request it just returns result, for cases where it's reasonable to retry it does it in
         exponential manner.
         """
+        # Clone kwargs and add Authorization header if auth is configured
+        request_kwargs = kwargs.copy()
+        if self.__auth:
+            auth_header = await self.__auth.get()
+            if auth_header:
+                if "headers" not in request_kwargs:
+                    request_kwargs["headers"] = {}
+                else:
+                    request_kwargs["headers"] = request_kwargs["headers"].copy()
+                request_kwargs["headers"]["Authorization"] = auth_header
+
         result = None
         exceptions = []
+
         for i in range(self.__retry_number + 1):  # add one for the first attempt, which is not a retry
             retry_factor = None
             try:
-                result = await method(url, **kwargs)
+                result = await method(url, **request_kwargs)
             except Exception as exc:
                 exceptions.append(exc)
                 if isinstance(exc, ServerConnectionError) or isinstance(exc, ClientResponseError):
@@ -104,6 +123,15 @@ async def __request(self, method: Callable, url, **kwargs: Any) -> ClientRespons
                     raise exc
 
             if result:
+                # Check for authentication errors first
+                if result.status in AUTH_PROBLEM_STATUSES and self.__auth:
+                    refreshed_header = await self.__auth.refresh()
+                    if refreshed_header:
+                        # Retry with new auth header
+                        request_kwargs["headers"] = request_kwargs.get("headers", {}).copy()
+                        request_kwargs["headers"]["Authorization"] = refreshed_header
+                        result = await method(url, **request_kwargs)
+
                 if result.ok or result.status not in RETRY_STATUSES:
                     return result
 
diff --git a/reportportal_client/_internal/services/auth.py b/reportportal_client/_internal/services/auth.py
@@ -29,9 +29,9 @@
 
 # noinspection PyAbstractClass
 class Auth(metaclass=AbstractBaseClass):
-    """Abstract base class for authentication.
+    """Abstract base class for synchronous authentication.
 
-    This class defines the interface for all authentication methods.
+    This class defines the interface for all synchronous authentication methods.
     """
 
     __metaclass__ = AbstractBaseClass
@@ -53,6 +53,32 @@ def refresh(self) -> Optional[str]:
         raise NotImplementedError('"refresh" method is not implemented!')
 
 
+# noinspection PyAbstractClass
+class AuthAsync(metaclass=AbstractBaseClass):
+    """Abstract base class for asynchronous authentication.
+
+    This class defines the interface for all asynchronous authentication methods.
+    """
+
+    __metaclass__ = AbstractBaseClass
+
+    @abstractmethod
+    async def get(self) -> Optional[str]:
+        """Get valid Authorization header value.
+
+        :return: Authorization header value or None if authentication failed.
+        """
+        raise NotImplementedError('"get" method is not implemented!')
+
+    @abstractmethod
+    async def refresh(self) -> Optional[str]:
+        """Refresh the access token and return Authorization header value.
+
+        :return: Authorization header value or None if refresh failed.
+        """
+        raise NotImplementedError('"refresh" method is not implemented!')
+
+
 class ApiKeyAuthSync(Auth):
     """Synchronous API key authentication.
 
@@ -86,7 +112,7 @@ def refresh(self) -> None:
         return None
 
 
-class ApiKeyAuthAsync(Auth):
+class ApiKeyAuthAsync(AuthAsync):
     """Asynchronous API key authentication.
 
     This class provides simple key-based authentication that always returns
@@ -120,15 +146,14 @@ async def refresh(self) -> None:
 
 
 # noinspection PyAbstractClass
-class OAuthPasswordGrant(Auth):
-    """Abstract base class for OAuth 2.0 password grant authentication.
+class OAuthPasswordGrant:
+    """Base class for OAuth 2.0 password grant authentication.
 
     This class provides common logic for obtaining and refreshing access tokens using
-    the OAuth 2.0 password grant flow.
+    the OAuth 2.0 password grant flow. This class should not be used directly, use
+    OAuthPasswordGrantSync or OAuthPasswordGrantAsync instead.
     """
 
-    __metaclass__ = AbstractBaseClass
-
     oauth_uri: str
     username: str
     password: str
@@ -244,7 +269,7 @@ def _build_token_request_data(self, grant_type: str, **extra_params) -> dict:
         return data
 
 
-class OAuthPasswordGrantSync(OAuthPasswordGrant):
+class OAuthPasswordGrantSync(OAuthPasswordGrant, Auth):
     """Synchronous implementation of OAuth 2.0 password grant authentication."""
 
     _session: Optional[requests.Session]
@@ -370,7 +395,7 @@ def close(self) -> None:
             self._session.close()
 
 
-class OAuthPasswordGrantAsync(OAuthPasswordGrant):
+class OAuthPasswordGrantAsync(OAuthPasswordGrant, AuthAsync):
     """Asynchronous implementation of OAuth 2.0 password grant authentication."""
 
     _session: Optional[aiohttp.ClientSession]
diff --git a/tests/_internal/aio/test_http.py b/tests/_internal/aio/test_http.py
@@ -31,6 +31,7 @@
 
 # noinspection PyProtectedMember
 from reportportal_client._internal.aio.http import RetryingClientSession
+from reportportal_client._internal.services.auth import ApiKeyAuthAsync
 
 HTTP_TIMEOUT_TIME = 1.2
 
@@ -75,6 +76,21 @@ def do_GET(self):
         self.wfile.flush()
 
 
+class UnauthorizedHttpHandler(http.server.BaseHTTPRequestHandler):
+    def do_GET(self):
+        auth_header = self.headers.get("Authorization")
+        if auth_header == "Bearer test_api_key":
+            self.send_response(200)
+            self.send_header("Content-Type", "application/json")
+            self.end_headers()
+            self.wfile.write("{}\n\n".encode("utf-8"))
+        else:
+            self.send_response(401, "Unauthorized")
+            self.end_headers()
+            self.wfile.write("Unauthorized\n\n".encode("utf-8"))
+        self.wfile.flush()
+
+
 SERVER_PORT = 8000
 SERVER_ADDRESS = ("", SERVER_PORT)
 SERVER_CLASS = socketserver.TCPServer
@@ -163,3 +179,93 @@ async def test_no_retry_on_not_retryable_error():
     assert result is None
     assert async_mock.call_count == 1
     assert total_time < 1
+
+
+@pytest.mark.asyncio
+async def test_auth_header_added_to_request():
+    """Test that auth header is added to requests when auth is configured."""
+    port = 8006
+    retry_number = 5
+    auth = ApiKeyAuthAsync("test_api_key")
+    timeout = aiohttp.ClientTimeout(connect=1.0, sock_read=1.0)
+    connector = aiohttp.TCPConnector(force_close=True)
+    session = RetryingClientSession(
+        f"http://localhost:{port}",
+        timeout=timeout,
+        max_retry_number=retry_number,
+        base_retry_delay=0.01,
+        auth=auth,
+        connector=connector,
+    )
+
+    with get_http_server(server_handler=UnauthorizedHttpHandler, server_address=("", port)):
+        async with session:
+            result = await session.get("/")
+            assert result.ok
+            assert result.status == 200
+
+
+@pytest.mark.asyncio
+async def test_auth_refresh_on_401():
+    """Test that 401 response triggers auth refresh."""
+    port = 8007
+    retry_number = 5
+
+    # Create a mock auth that fails first, then succeeds
+    auth = mock.AsyncMock()
+    auth.get = mock.AsyncMock(side_effect=["Bearer invalid_token", "Bearer test_api_key"])
+    auth.refresh = mock.AsyncMock(return_value="Bearer test_api_key")
+
+    timeout = aiohttp.ClientTimeout(connect=1.0, sock_read=1.0)
+    connector = aiohttp.TCPConnector(force_close=True)
+    session = RetryingClientSession(
+        f"http://localhost:{port}",
+        timeout=timeout,
+        max_retry_number=retry_number,
+        base_retry_delay=0.01,
+        auth=auth,
+        connector=connector,
+    )
+
+    with get_http_server(server_handler=UnauthorizedHttpHandler, server_address=("", port)):
+        async with session:
+            result = await session.get("/")
+            # First call to get() returns invalid token, which causes 401
+            # Then refresh() is called and returns valid token
+            # Request is retried with valid token and succeeds
+            assert result.ok
+            assert result.status == 200
+            assert auth.get.call_count == 1
+            assert auth.refresh.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_auth_refresh_only_once():
+    """Test that auth refresh is only performed once per request."""
+    port = 8008
+    retry_number = 5
+
+    # Create a mock auth that always fails
+    auth = mock.AsyncMock()
+    auth.get = mock.AsyncMock(return_value="Bearer invalid_token")
+    auth.refresh = mock.AsyncMock(return_value="Bearer still_invalid_token")
+
+    timeout = aiohttp.ClientTimeout(connect=1.0, sock_read=1.0)
+    connector = aiohttp.TCPConnector(force_close=True)
+    session = RetryingClientSession(
+        f"http://localhost:{port}",
+        timeout=timeout,
+        max_retry_number=retry_number,
+        base_retry_delay=0.01,
+        auth=auth,
+        connector=connector,
+    )
+
+    with get_http_server(server_handler=UnauthorizedHttpHandler, server_address=("", port)):
+        async with session:
+            result = await session.get("/")
+            # Auth refresh should only be attempted once
+            assert not result.ok
+            assert result.status == 401
+            assert auth.get.call_count == 1
+            assert auth.refresh.call_count == 1
diff --git a/tests/_internal/services/test_auth.py b/tests/_internal/services/test_auth.py
@@ -542,26 +542,24 @@ def test_get_returns_token(self):
 
         assert result == f"Bearer {api_token}"
 
-    def test_refresh_returns_token(self):
-        """Test that refresh() returns the API token."""
+    def test_refresh_returns_none(self):
+        """Test that refresh() returns None (API keys don't have refresh mechanism)."""
         api_token = "test_api_token_67890"
         auth = ApiKeyAuthSync(api_token)
         result = auth.refresh()
 
-        assert result == f"Bearer {api_token}"
+        assert result is None
 
     def test_multiple_calls_return_same_token(self):
-        """Test that multiple calls return the same token."""
+        """Test that multiple get() calls return the same token."""
         api_token = "test_api_token_stable"
         auth = ApiKeyAuthSync(api_token)
 
         result1 = auth.get()
         result2 = auth.get()
-        result3 = auth.refresh()
 
         assert result1 == f"Bearer {api_token}"
         assert result2 == f"Bearer {api_token}"
-        assert result3 == f"Bearer {api_token}"
 
 
 class TestApiTokenAuthAsync:
@@ -577,24 +575,22 @@ async def test_get_returns_token(self):
         assert result == f"Bearer {api_token}"
 
     @pytest.mark.asyncio
-    async def test_refresh_returns_token(self):
-        """Test that refresh() returns the API token."""
+    async def test_refresh_returns_none(self):
+        """Test that refresh() returns None (API keys don't have refresh mechanism)."""
         api_token = "test_api_token_async_67890"
         auth = ApiKeyAuthAsync(api_token)
         result = await auth.refresh()
 
-        assert result == f"Bearer {api_token}"
+        assert result is None
 
     @pytest.mark.asyncio
     async def test_multiple_calls_return_same_token(self):
-        """Test that multiple calls return the same token."""
+        """Test that multiple get() calls return the same token."""
         api_token = "test_api_token_async_stable"
         auth = ApiKeyAuthAsync(api_token)
 
         result1 = await auth.get()
         result2 = await auth.get()
-        result3 = await auth.refresh()
 
         assert result1 == f"Bearer {api_token}"
         assert result2 == f"Bearer {api_token}"
-        assert result3 == f"Bearer {api_token}"
