0.0.165

orso/logging/google_cloud_logger.py

@@ -68,9 +68,7 @@ def extract_caller():
 class GoogleLogger(object):
     @staticmethod
     def supported():
-        if os.environ.get("K_SERVICE", "") or "" != "":  # nosemgrep
-            return True
-        return False
+        return os.environ.get("K_SERVICE", "") or "" != ""  # nosemgrep
 
     @staticmethod
     def write_event(

orso/tools.py

@@ -12,7 +12,9 @@
 
 import datetime
 import decimal
+import logging
 import os
+import random
 import threading
 import time
 import uuid
@@ -21,22 +23,29 @@
 from typing import Any
 from typing import Callable
 from typing import List
+from typing import Optional
+from typing import Tuple
 from typing import Type
 from typing import Union
 
 import numpy
 
 from orso.exceptions import MissingDependencyError
 
+logger = logging.getLogger(__name__)
+
 
 def retry(
     max_tries: int = 3,
     backoff_seconds: int = 1,
     exponential_backoff: bool = False,
     max_backoff: int = 4,
+    retry_exceptions: Tuple[Type[Exception]] = (Exception,),
+    jitter: bool = False,
+    callback: Optional[Callable[[Exception, int], None]] = None,
 ) -> Callable:
     """
-    Decorator to add retry logic with optional exponential backoff to a function.
+    Decorator to add retry logic with optional exponential backoff and jitter to a function.
 
     Parameters:
         max_tries: int
@@ -47,6 +56,12 @@ def retry(
             Whether to use exponential backoff for the delay between retries.
         max_backoff: int
             The maximum backoff time (in seconds) when using exponential backoff.
+        retry_exceptions: Tuple[Type[Exception]]
+            Tuple of exception types that should trigger a retry.
+        jitter: bool
+            Whether to add a small random delay to the backoff time.
+        callback: Optional[Callable[[Exception, int], None]]
+            A callback function to be called after each failure. It receives the exception and the attempt number.
 
     Returns:
         Callable: Wrapped function with retry logic.
@@ -58,30 +73,30 @@ def wrapper_retry(*args, **kwargs):
             tries = 0
             this_delay = backoff_seconds
 
-            # Retry logic
             while tries < max_tries:
                 try:
                     return func(*args, **kwargs)
-                except Exception as e:
+                except retry_exceptions as e:
                     tries += 1
 
-                    # Reached maximum retries, raise the exception
+                    if callback:
+                        callback(e, tries)
+
                     if tries == max_tries:
-                        print(
-                            f"`{func.__name__}` failed with `{type(e).__name__}` error after {tries} attempts. Aborting."
+                        logger.error(
+                            f"`{func.__name__}` failed with `{type(e).__name__}` after {tries} attempts. Aborting."
                         )
                         raise e
 
-                    # Log the exception and retry after waiting
-                    print(
-                        f"`{func.__name__}` failed with `{type(e).__name__}` error, attempt {tries} of {max_tries}. Will retry in {this_delay} seconds."
+                    logger.warning(
+                        f"`{func.__name__}` failed with `{type(e).__name__}` error, attempt {tries} of {max_tries}. Retrying in {this_delay} seconds."
                     )
                     time.sleep(this_delay)
 
-                    # Update delay time for exponential backoff
                     if exponential_backoff:
-                        this_delay *= 2
-                        this_delay = min(this_delay, max_backoff)
+                        this_delay = min(this_delay * 2, max_backoff)
+                    if jitter:
+                        this_delay += random.uniform(0, 0.5)
 
         return wrapper_retry
 

orso/version.py

@@ -10,5 +10,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__: str = "0.0.164"
+__version__: str = "0.0.165"
 __author__: str = "@joocer"

tests/test_tools_retry.py

@@ -0,0 +1,177 @@
+import os
+import sys
+import pytest
+
+sys.path.insert(1, os.path.join(sys.path[0], ".."))
+
+import pytest
+from unittest.mock import Mock, patch
+from orso.tools import retry  # assuming the decorator is saved in retry_decorator.py
+
+
+def test_retry_success_on_first_try():
+    mock_func = Mock(return_value="success")
+
+    @retry(max_tries=3)
+    def decorated_func():
+        return mock_func()
+
+    with patch('time.sleep', return_value=None) as mock_sleep:
+        result = decorated_func()
+        assert result == "success"
+        mock_func.assert_called_once()
+        mock_sleep.assert_not_called()
+
+def test_retry_success_after_retries():
+    mock_func = Mock(side_effect=[Exception("fail"), Exception("fail"), "success"])
+
+    @retry(max_tries=3, backoff_seconds=1)
+    def decorated_func():
+        return mock_func()
+
+    with patch('time.sleep', return_value=None) as mock_sleep:
+        result = decorated_func()
+        assert result == "success"
+        assert mock_func.call_count == 3
+        assert mock_sleep.call_count == 2
+
+def test_retry_failure_after_max_retries():
+    mock_func = Mock(side_effect=Exception("fail"))
+
+    @retry(max_tries=3, backoff_seconds=1)
+    def decorated_func():
+        return mock_func()
+
+    with patch('time.sleep', return_value=None) as mock_sleep:
+        with pytest.raises(Exception):
+            decorated_func()
+
+        assert mock_func.call_count == 3
+        assert mock_sleep.call_count == 2
+
+def test_retry_with_exponential_backoff():
+    mock_func = Mock(side_effect=[Exception("fail"), Exception("fail"), "success"])
+
+    @retry(max_tries=3, backoff_seconds=1, exponential_backoff=True, max_backoff=4)
+    def decorated_func():
+        return mock_func()
+
+    with patch('time.sleep', return_value=None) as mock_sleep:
+        result = decorated_func()
+        assert result == "success"
+        assert mock_func.call_count == 3
+        assert mock_sleep.call_count == 2
+        mock_sleep.assert_any_call(1)
+        mock_sleep.assert_any_call(2)
+
+def test_retry_with_jitter():
+    mock_func = Mock(side_effect=[Exception("fail"), Exception("fail"), "success"])
+
+    @retry(max_tries=3, backoff_seconds=1, jitter=True)
+    def decorated_func():
+        return mock_func()
+
+    with patch('time.sleep', return_value=None) as mock_sleep, patch('random.uniform', return_value=0.5) as mock_random:
+        result = decorated_func()
+        assert result == "success"
+        assert mock_func.call_count == 3
+        assert mock_sleep.call_count == 2
+        mock_random.assert_called()
+
+def test_retry_with_specific_exceptions():
+    mock_func = Mock(side_effect=[ValueError("fail"), ValueError("fail"), "success"])
+
+    @retry(max_tries=3, backoff_seconds=1, retry_exceptions=(ValueError,))
+    def decorated_func():
+        return mock_func()
+
+    with patch('time.sleep', return_value=None) as mock_sleep:
+        result = decorated_func()
+        assert result == "success"
+        assert mock_func.call_count == 3
+        assert mock_sleep.call_count == 2
+
+def test_retry_with_callback():
+    mock_func = Mock(side_effect=[Exception("fail"), "success"])
+    callback_func = Mock()
+
+    @retry(max_tries=3, backoff_seconds=1, callback=callback_func)
+    def decorated_func():
+        return mock_func()
+
+    with patch('time.sleep', return_value=None) as mock_sleep:
+        result = decorated_func()
+        assert result == "success"
+        assert mock_func.call_count == 2
+        assert mock_sleep.call_count == 1
+
+        # Ensure callback was called once with correct arguments
+        assert callback_func.call_count == 1
+        called_exception, called_attempt = callback_func.call_args[0]
+        assert isinstance(called_exception, Exception)
+        assert called_exception.args == ("fail",)
+        assert called_attempt == 1
+
+
+def test_retry_no_retries_needed():
+    mock_func = Mock(return_value="success")
+
+    @retry(max_tries=3, backoff_seconds=1)
+    def decorated_func():
+        return mock_func()
+
+    with patch('time.sleep', return_value=None) as mock_sleep:
+        result = decorated_func()
+        assert result == "success"
+        assert mock_func.call_count == 1
+        mock_sleep.assert_not_called()
+
+def test_retry_with_multiple_exceptions():
+    mock_func = Mock(side_effect=[ValueError("fail"), KeyError("fail"), "success"])
+
+    @retry(max_tries=3, backoff_seconds=1, retry_exceptions=(ValueError, KeyError))
+    def decorated_func():
+        return mock_func()
+
+    with patch('time.sleep', return_value=None) as mock_sleep:
+        result = decorated_func()
+        assert result == "success"
+        assert mock_func.call_count == 3
+        assert mock_sleep.call_count == 2
+
+def test_retry_with_partial_success():
+    mock_func = Mock(side_effect=[Exception("fail"), "partial_success", Exception("fail_again"), "success"])
+
+    @retry(max_tries=4, backoff_seconds=1)
+    def decorated_func():
+        result = mock_func()
+        if result == "partial_success":
+            raise Exception("retrying due to partial success")
+        return result
+
+    with patch('time.sleep', return_value=None) as mock_sleep:
+        result = decorated_func()
+        assert result == "success"
+        assert mock_func.call_count == 4
+        assert mock_sleep.call_count == 3
+
+def test_retry_with_max_backoff():
+    mock_func = Mock(side_effect=[Exception("fail"), Exception("fail"), "success"])
+
+    @retry(max_tries=3, backoff_seconds=1, exponential_backoff=True, max_backoff=2)
+    def decorated_func():
+        return mock_func()
+
+    with patch('time.sleep', return_value=None) as mock_sleep:
+        result = decorated_func()
+        assert result == "success"
+        assert mock_func.call_count == 3
+        assert mock_sleep.call_count == 2
+        mock_sleep.assert_any_call(1)
+        mock_sleep.assert_any_call(2)
+
+
+if __name__ == "__main__":  # pragma: nocover
+    from tests import run_tests
+
+    run_tests()