Python Typed Parameterized Decorators in Test Automation

Nikolai Filatov — Wed, 22 Jan 2025 11:07:51 +0000

Python's decorator functionality, enriched with the latest type hinting features, transforms how we approach test automation. By leveraging Python's flexibility with decorators and the enhanced type safety provided by the typing module, we can create more maintainable, readable, and robust test suites. This article delves into how to use these advanced features, focusing on their application within test automation frameworks.

Latest Features from the Typing Module

Python's typing module has seen significant updates:

PEP 585 introduces native support for generic types in standard collections, reducing the use of the typing module for common types
PEP 604 introduces the | operator for Union types, simplifying type annotations
PEP 647 introduces TypeAlias for clearer type alias definitions
PEP 649 allows for deferred evaluation of annotations, improving startup time for large projects

Creating Typed Parameterized Decorators

Here's how to craft a decorator using these new typing features:

from typing import Protocol, TypeVar, Generic, Callable, Any
from functools import wraps

# TypeVar for generic typing
T = TypeVar('T')

# Protocol for defining function structure
class TestProtocol(Protocol):
    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        ...

def generic_decorator(param: str) -> Callable[[Callable[..., T]], Callable[..., T]]:
    """
    A generic decorator that can wrap any function returning type T.

    Args:
        param: A string parameter used by the decorator

    Returns:
        A callable that wraps the original function
    """
    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        @wraps(func)  # Preserves the metadata of the original function
        def wrapper(*args: Any, **kwargs: Any) -> T:
            print(f"Generic decorator with param: {param}")
            return func(*args, **kwargs)
        return wrapper
    return decorator

@generic_decorator("test_param")
def test_function(x: int) -> int:
    """A test function that returns the input multiplied by 2."""
    return x * 2

This decorator uses Protocol to define what a test function should look like, enhancing flexibility when dealing with different function signatures in test frameworks.

Applying to Test Automation Frameworks

Let's explore how these decorators can be used in test automation:

1. Platform-Specific Tests with Literal

from typing import Literal, Callable, Any
import sys

def run_only_on(platform: Literal["linux", "darwin", "win32"]) -> Callable:
    """
    Decorator to ensure a test function runs only on specified platforms.

    Args:
        platform: The target platform for the test

    Returns:
        A callable that wraps the original test function
    """
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            if sys.platform == platform:
                return func(*args, **kwargs)
            print(f"Skipping test for platform: {sys.platform}")
            return None
        return wrapper
    return decorator

@run_only_on("linux")
def test_linux_feature() -> None:
    """Linux-specific test."""
    # Linux-specific operations
    pass

The Literal type ensures type checkers know exactly which values are valid for the platform parameter, making it clear which tests run on which platforms. This is particularly useful in cross-platform testing frameworks.

2. Timeout Decorators with Threading

from typing import Callable, Any, Optional
import threading
import time
from concurrent.futures import ThreadPoolExecutor, TimeoutError

def timeout(seconds: int) -> Callable:
    """
    Decorator to enforce a timeout on test functions.

    Args:
        seconds: Maximum execution time in seconds

    Returns:
        A callable that wraps the original function with timeout functionality
    """
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Optional[Any]:
            with ThreadPoolExecutor(max_workers=1) as executor:
                future = executor.submit(func, *args, **kwargs)
                try:
                    return future.result(timeout=seconds)
                except TimeoutError:
                    print(f"Function {func.__name__} timed out after {seconds} seconds")
                    return None
        return wrapper
    return decorator

@timeout(5)
def test_long_running_operation() -> None:
    """A test that times out if it takes too long."""
    time.sleep(10)  # This will trigger a timeout

This implementation uses Python's built-in threading capabilities to implement a reliable timeout mechanism, suitable for test automation scenarios where execution time must be constrained.

3. Retry Mechanism with Union Types

from typing import Callable, Any, Union, Type, Tuple, Optional
import time

def retry_on_exception(
    exceptions: Union[Type[Exception], Tuple[Type[Exception], ...]], 
    attempts: int = 3,
    delay: float = 1.0
) -> Callable:
    """
    Decorator to retry a function on specified exceptions.

    Args:
        exceptions: Single exception or tuple of exceptions to catch
        attempts: Maximum number of retry attempts
        delay: Delay between attempts in seconds

    Returns:
        A callable that wraps the original function with retry logic
    """
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            last_exception: Optional[Exception] = None

            for attempt in range(attempts):
                try:
                    return func(*args, **kwargs)
                except exceptions as e:
                    last_exception = e
                    print(f"Attempt {attempt + 1} failed with {type(e).__name__}: {str(e)}")
                    if attempt < attempts - 1:
                        time.sleep(delay)

            if last_exception is not None:
                raise last_exception
            return None
        return wrapper
    return decorator

@retry_on_exception((TimeoutError, ConnectionError), attempts=5, delay=1.0)
def test_network_connection() -> None:
    """Test network connection with retry logic."""
    # Simulate network operations
    pass

This improved version includes comprehensive type hints, robust exception handling, and a configurable delay between retries. It uses Union types to accept either a single exception type or a tuple of exception types, providing flexibility for various testing scenarios.

Conclusion

The integration of Python's latest typing features into decorators enhances both type safety and code readability while improving test automation framework capabilities. By defining explicit behavior through types, you ensure that tests execute under the right conditions with proper error handling and performance constraints. This approach leads to more robust, maintainable, and efficient testing processes, particularly beneficial in large-scale, distributed, or multi-platform test environments.

DEV Community: Nikolai Filatov