Update retries settings logic and add tests

reportportal · Oct 2, 2023 · c2af6b0 · c2af6b0
1 parent c62a571
commit c2af6b0
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 4 deletions.
diff --git a/reportportal_client/aio/client.py b/reportportal_client/aio/client.py
@@ -45,7 +45,7 @@
     AbstractBaseClass,
     abstractmethod
 )
-from reportportal_client.static.defines import NOT_FOUND
+from reportportal_client.static.defines import NOT_FOUND, NOT_SET
 from reportportal_client.steps import StepReporter
 
 logger = logging.getLogger(__name__)
@@ -93,7 +93,7 @@ def __init__(
             api_key: str = None,
             is_skipped_an_issue: bool = True,
             verify_ssl: Union[bool, str] = True,
-            retries: int = None,
+            retries: int = NOT_SET,
             max_pool_size: int = 50,
             http_timeout: Union[float, Tuple[float, float]] = (10, 10),
             keepalive_timeout: Optional[float] = None,
@@ -198,10 +198,16 @@ def session(self) -> aiohttp.ClientSession:
                 connect_timeout, read_timeout = self.http_timeout, self.http_timeout
             session_params['timeout'] = aiohttp.ClientTimeout(connect=connect_timeout, sock_read=read_timeout)
 
-        if self.retries:
+        retries_set = self.retries is not NOT_SET and self.retries and self.retries > 0
+        use_retries = self.retries is NOT_SET or (self.retries and self.retries > 0)
+
+        if retries_set:
             session_params['max_retry_number'] = self.retries
 
-        self._session = RetryingClientSession(self.endpoint, **session_params)
+        if use_retries:
+            self._session = RetryingClientSession(self.endpoint, **session_params)
+        else:
+            self._session = aiohttp.ClientSession(self.endpoint, **session_params)
         return self._session
 
     async def close(self) -> None:

diff --git a/tests/aio/test_aio_client.py b/tests/aio/test_aio_client.py
@@ -12,12 +12,37 @@
 #  limitations under the License
 
 import pickle
+from unittest import mock
+
+import aiohttp
+import pytest
 
 from reportportal_client.aio.client import Client
+from reportportal_client.aio.http import RetryingClientSession, DEFAULT_RETRY_NUMBER
+from reportportal_client.static.defines import NOT_SET
 
 
 def test_client_pickling():
     client = Client('http://localhost:8080', 'default_personal', api_key='test_key')
     pickled_client = pickle.dumps(client)
     unpickled_client = pickle.loads(pickled_client)
     assert unpickled_client is not None
+
+
+@pytest.mark.parametrize(
+    'retry_num, expected_class, expected_param',
+    [
+        (1, RetryingClientSession, 1),
+        (0, aiohttp.ClientSession, NOT_SET),
+        (-1, aiohttp.ClientSession, NOT_SET),
+        (None, aiohttp.ClientSession, NOT_SET),
+        (NOT_SET, RetryingClientSession, DEFAULT_RETRY_NUMBER)
+    ]
+)
+def test_retries_param(retry_num, expected_class, expected_param):
+    client = Client('http://localhost:8080', 'default_personal', api_key='test_key',
+                    retries=retry_num)
+    session = client.session
+    assert isinstance(session, expected_class)
+    if expected_param is not NOT_SET:
+        assert getattr(session, f'_RetryingClientSession__retry_number') == expected_param