Retry Logic and Exception Handling

def ensure(*, attempts: Optional[int] = None,
           delay: Optional[Union[float, int, Callable[[int], Union[float, int]]]] = None,
           swallow: Optional[Union[Type[BaseException], Tuple[Type[BaseException], ...]]] = None) -> Ensure:
    """
    Decorator to add retry logic to a function or coroutine.
    
    Args:
        attempts: The maximum number of times the function can be called.
                 To run the function and retry once if an exception is raised, set attempts=2.
                 To run the function and retry twice, set attempts=3.
        delay: The delay between attempts, which can be a fixed value or a callable
               returning a value.
        swallow: The exception(s) to be caught and retried.
        
    Returns:
        An Ensure instance configured with the provided or default parameters.
    """

class Ensure:
    """
    Provides functionality to ensure a function succeeds within a specified number of attempts.
    
    This class retries a given function or coroutine function a specified number of times,
    optionally with a delay between attempts, and can log each attempt.
    """
    
    def __init__(self, *, attempts: int = 3,
                 delay: Union[float, int, Callable[[int], Union[float, int]]] = 0.0,
                 swallow: Union[Type[BaseException], Tuple[Type[BaseException], ...]] = BaseException,
                 logger: Optional[Callable] = None): ...
    
    def __call__(self, fn: Callable) -> Callable: ...

from vedro import scenario, given, when, then, ensure

@scenario("Flaky API operation")
def test_api_with_retries():
    
    @given("API client")
    def setup():
        return APIClient("https://flaky-service.com")
    
    @when("making request with retry logic")
    def action(client):
        # Retry up to 3 times with 1 second delay
        @ensure(attempts=3, delay=1.0, swallow=(ConnectionError, TimeoutError))
        def make_request():
            return client.get("/users/123")
        
        return make_request()
    
    @then("request eventually succeeds")
    def verification(response):
        assert response.status_code == 200
        assert "user_id" in response.json()

class catched:
    """
    Context manager for catching and inspecting exceptions.
    
    Supports both synchronous and asynchronous contexts.
    Can be used to verify that specific exceptions are raised
    and inspect their properties.
    """
    
    def __init__(self, expected_exc = BaseException):
        """
        Initialize exception catcher.
        
        Args:
            expected_exc: Exception type(s) to catch. Can be single type
                         or tuple of types. Defaults to BaseException.
        """
    
    @property  
    def type(self) -> Type[BaseException] | None:
        """The type of the caught exception, if any."""
    
    @property
    def value(self) -> BaseException | None: 
        """The caught exception instance, if any."""
    
    @property
    def traceback(self) -> TracebackType | None:
        """The traceback of the caught exception, if any."""
    
    def __enter__(self) -> "catched": ...
    def __exit__(self, exc_type, exc_value, traceback) -> bool: ...
    def __aenter__(self) -> "catched": ...
    def __aexit__(self, exc_type, exc_value, traceback) -> bool: ...
    def __repr__(self) -> str: ...

from vedro import scenario, given, when, then, catched

@scenario("Invalid input handling")
def test_invalid_input():
    
    @given("invalid user data")
    def setup():
        return {
            "name": "",  # Invalid: empty name
            "email": "not-an-email",  # Invalid: malformed email
            "age": -5  # Invalid: negative age
        }
    
    @when("attempting to create user")
    def action(invalid_data):
        with catched(ValidationError) as caught:
            create_user(invalid_data)
        return caught
    
    @then("appropriate validation error is raised")
    def verification(caught_exception):
        # Verify exception was caught
        assert caught_exception.value is not None
        assert caught_exception.type == ValidationError
        
        # Inspect exception details
        error = caught_exception.value
        assert "name" in error.field_errors
        assert "email" in error.field_errors
        assert "age" in error.field_errors
        assert "validation failed" in str(error)

@scenario("Database operation error handling")
def test_database_errors():
    
    @when("database operation fails")
    def action():
        # Catch multiple possible exception types
        with catched((ConnectionError, TimeoutError, DatabaseError)) as caught:
            unreliable_database_operation()
        return caught
    
    @then("appropriate error handling occurs")
    def verification(caught_exception):
        assert caught_exception.value is not None
        
        # Handle different exception types appropriately
        if caught_exception.type == ConnectionError:
            assert "connection" in str(caught_exception.value)
        elif caught_exception.type == TimeoutError:
            assert "timeout" in str(caught_exception.value)
        elif caught_exception.type == DatabaseError:
            assert "database" in str(caught_exception.value)

@scenario("Async operation error handling")
def test_async_errors():
    
    @when("async operation fails")
    async def action():
        async with catched(AsyncOperationError) as caught:
            await failing_async_operation()
        return caught
    
    @then("async error is properly caught")
    def verification(caught_exception):
        assert caught_exception.value is not None
        assert caught_exception.type == AsyncOperationError
        
        # Check async-specific error details
        error = caught_exception.value
        assert error.operation_id is not None
        assert error.retry_count > 0

@scenario("Resilient API operations")
def test_resilient_operations():
    
    @when("performing operation with retries and error handling")
    def action():
        @ensure(attempts=3, delay=0.5, swallow=(ConnectionError, TimeoutError))
        def reliable_operation():
            # This might fail with connection issues but will be retried
            return api_call_that_might_fail()
        
        try:
            result = reliable_operation()
            return {"success": True, "result": result}
        except Exception as e:
            with catched(type(e)) as caught:
                raise
            return {"success": False, "caught": caught}
    
    @then("operation handles failures gracefully")
    def verification(outcome):
        if outcome["success"]:
            assert outcome["result"] is not None
        else:
            caught = outcome["caught"]
            assert caught.value is not None
            # Verify it's not a retryable error (those should have been handled)
            assert not isinstance(caught.value, (ConnectionError, TimeoutError))

from types import TracebackType
from typing import Type, Union, Tuple, Callable, Optional

# Retry mechanism types
AttemptType = int
DelayValueType = Union[float, int]
DelayCallableType = Callable[[AttemptType], DelayValueType]
DelayType = Union[DelayValueType, DelayCallableType]

ExceptionType = Type[BaseException]
SwallowExceptionType = Union[Tuple[ExceptionType, ...], ExceptionType]

LoggerType = Callable[[Callable, AttemptType, Union[BaseException, None]], Any]

# Exception handling types
ExpectedExcType = Union[Type[BaseException], Tuple[Type[BaseException], ...]]

@scenario("Exponential backoff retry")
def test_exponential_backoff():
    
    @when("using exponential backoff for unreliable service")
    def action():
        def exponential_delay(attempt: int) -> float:
            return min(2 ** attempt, 30)  # Cap at 30 seconds
        
        @ensure(attempts=5, delay=exponential_delay, swallow=ServiceUnavailableError)
        def call_unreliable_service():
            return external_api_call()
        
        return call_unreliable_service()
    
    @then("service call eventually succeeds with backoff")
    def verification(result):
        assert result is not None
        assert result.status == "success"

@scenario("Exception chaining")
def test_exception_chain():
    
    @when("nested operation fails")
    def action():
        with catched(ServiceError) as caught:
            # This should catch a ServiceError caused by a DatabaseError
            complex_service_operation()
        return caught
    
    @then("exception chain is preserved")
    def verification(caught_exception):
        service_error = caught_exception.value
        assert service_error is not None
        assert service_error.__cause__ is not None
        assert type(service_error.__cause__) == DatabaseError
        
        # Verify the full chain
        db_error = service_error.__cause__
        assert db_error.connection_string is not None
        assert "service unavailable" in str(service_error)
        assert "connection failed" in str(db_error)