diff --git a/robust/exception.py b/robust/exception.py index e72ad04..458e7ef 100644 --- a/robust/exception.py +++ b/robust/exception.py @@ -4,3 +4,7 @@ class ContinuousFailureException(Exception): class TimeoutException(Exception): pass + + +class ConnectionCutException(Exception): + pass diff --git a/robust/tests.py b/robust/tests.py index 0306f8b..a233418 100644 --- a/robust/tests.py +++ b/robust/tests.py @@ -1,6 +1,8 @@ from pytest import raises from robust.tools import retry, timeout, breaker -from robust.exception import ContinuousFailureException, TimeoutException +from robust.exception import (ContinuousFailureException, + TimeoutException, + ConnectionCutException) class TestRetryCase(object): diff --git a/robust/tools.py b/robust/tools.py index 245d85d..54d3b7d 100644 --- a/robust/tools.py +++ b/robust/tools.py @@ -59,3 +59,22 @@ def breaker(limit, revive, on_fail=None): After :revive: seconds it allows one connection to pass. If it succeeds - counter is reset, if doesn't - we wait another :revive: seconds """ + + def injector(fn): + counter = 0 + + @wraps(fn) + def wrapper(*args, **kwargs): + nonlocal counter + if counter > limit: + pass + + try: + return fn(*args, **kwargs) + except Exception: + counter += 1 + raise + + return wrapper + + return injector