frostfs-testlib/src/frostfs_testlib/testing/test_control.py

import inspect
import logging
from functools import wraps
from time import sleep, time
from typing import Any

from _pytest.outcomes import Failed
from pytest import fail

logger = logging.getLogger("NeoLogger")

# TODO: we may consider deprecating some methods here and use tenacity instead


class expect_not_raises:
    """
    Decorator/Context manager check that some action, method or test does not raise exceptions

    Useful to set proper state of failed test cases in allure

    Example:
        def do_stuff():
            raise Exception("Fail")

        def test_yellow(): <- this test is marked yellow (Test Defect) in allure
            do_stuff()

        def test_red(): <- this test is marked red (Failed) in allure
            with expect_not_raises():
                do_stuff()

        @expect_not_raises()
        def test_also_red(): <- this test is also marked red (Failed) in allure
            do_stuff()
    """

    def __enter__(self):
        pass

    def __exit__(self, exception_type, exception_value, exception_traceback):
        if exception_value:
            fail(str(exception_value))

    def __call__(self, func):
        @wraps(func)
        def impl(*a, **kw):
            with expect_not_raises():
                func(*a, **kw)

        return impl


def retry(max_attempts: int, sleep_interval: int = 1, expected_result: Any = None):
    """
    Decorator to wait for some conditions/functions to pass successfully.
    This is useful if you don't know exact time when something should pass successfully and do not
    want to use sleep(X) with too big X.

    Be careful though, wrapped function should only check the state of something, not change it.
    """

    assert max_attempts >= 1, "Cannot apply retry decorator with max_attempts < 1"

    def wrapper(func):
        @wraps(func)
        def impl(*a, **kw):
            last_exception = None
            for _ in range(max_attempts):
                try:
                    actual_result = func(*a, **kw)
                    if expected_result is not None:
                        assert expected_result == actual_result
                    return actual_result
                except Exception as ex:
                    logger.debug(ex)
                    last_exception = ex
                    sleep(sleep_interval)
                except Failed as ex:
                    logger.debug(ex)
                    last_exception = ex
                    sleep(sleep_interval)

            # timeout exceeded with no success, raise last_exception
            if last_exception is not None:
                raise last_exception

        return impl

    return wrapper


def run_optionally(enabled: bool, mock_value: Any = True):
    """
    Decorator to run something conditionally.
    MUST be placed after @pytest.fixture and before @allure decorators.

    Args:
    enabled: if true, decorated func will be called as usual. if false the decorated func will be skipped and mock_value will be returned.
    mock_value: the value to be returned when decorated func is skipped.
    """

    def deco(func):
        @wraps(func)
        def func_impl(*a, **kw):
            if enabled:
                return func(*a, **kw)
            return mock_value

        @wraps(func)
        def gen_impl(*a, **kw):
            if enabled:
                yield from func(*a, **kw)
                return
            yield mock_value

        return gen_impl if inspect.isgeneratorfunction(func) else func_impl

    return deco


def wait_for_success(
    max_wait_time: int = 60,
    interval: int = 1,
    expected_result: Any = None,
    fail_testcase: bool = False,
    fail_message: str = "",
):
    """
    Decorator to wait for some conditions/functions to pass successfully.
    This is useful if you don't know exact time when something should pass successfully and do not
    want to use sleep(X) with too big X.

    Be careful though, wrapped function should only check the state of something, not change it.
    """

    def wrapper(func):
        @wraps(func)
        def impl(*a, **kw):
            start = int(round(time()))
            last_exception = None
            while start + max_wait_time >= int(round(time())):
                try:
                    actual_result = func(*a, **kw)
                    if expected_result is not None:
                        assert expected_result == actual_result, fail_message
                    return actual_result
                except Exception as ex:
                    logger.debug(ex)
                    last_exception = ex
                    sleep(interval)
                except Failed as ex:
                    logger.debug(ex)
                    last_exception = ex
                    sleep(interval)

            if fail_testcase:
                fail(str(last_exception))

            # timeout exceeded with no success, raise last_exception
            if last_exception is not None:
                raise last_exception

        return impl

    return wrapper