Merge pull request #16 from AugPro/develop

Adding support for generator functions and coroutines

circuitbreaker.py

@@ -6,6 +6,7 @@
 
 from functools import wraps
 from datetime import datetime, timedelta
+from inspect import isgeneratorfunction
 from typing import AnyStr, Iterable
 
 STATE_CLOSED = 'closed'
@@ -38,6 +39,18 @@ def __init__(self,
     def __call__(self, wrapped):
         return self.decorate(wrapped)
 
+    def __enter__(self):
+        return None
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type and issubclass(exc_type, self._expected_exception):
+            # exception was raised and is our concern
+            self._last_failure = exc_value
+            self.__call_failed()
+        else:
+            self.__call_succeeded()
+        return False  # return False to raise exception if any
+
     def decorate(self, function):
         """
         Applies the circuit breaker to a function
@@ -47,9 +60,18 @@ def decorate(self, function):
 
         CircuitBreakerMonitor.register(self)
 
+        if isgeneratorfunction(function):
+            call = self.call_generator
+        else:
+            call = self.call
+
         @wraps(function)
         def wrapper(*args, **kwargs):
-            return self.call(function, *args, **kwargs)
+            if self.opened:
+                if self.fallback_function:
+                    return self.fallback_function(*args, **kwargs)
+                raise CircuitBreakerError(self)
+            return call(function, *args, **kwargs)
 
         return wrapper
 
@@ -59,19 +81,18 @@ def call(self, func, *args, **kwargs):
         rules on success or failure
         :param func: Decorated function
         """
-        if self.opened:
-            if self.fallback_function:
-                return self.fallback_function(*args, **kwargs)
-            raise CircuitBreakerError(self)
-        try:
-            result = func(*args, **kwargs)
-        except self._expected_exception as e:
-            self._last_failure = e
-            self.__call_failed()
-            raise
+        with self:
+            return func(*args, **kwargs)
 
-        self.__call_succeeded()
-        return result
+    def call_generator(self, func, *args, **kwargs):
+        """
+        Calls the decorated generator function and applies the circuit breaker
+        rules on success or failure
+        :param func: Decorated genrator function
+        """
+        with self:
+            for el in func(*args, **kwargs):
+                yield el
 
     def __call_succeeded(self):
         """

tests/test_functional.py

@@ -21,6 +21,13 @@ def circuit_failure():
     raise IOError()
 
 
+@CircuitBreaker(failure_threshold=1, name="circuit_generator_failure")
+def circuit_generator_failure():
+    pseudo_remote_call()
+    yield 1
+    raise IOError()
+
+
 @CircuitBreaker(failure_threshold=1, name="threshold_1")
 def circuit_threshold_1():
     return pseudo_remote_call()
@@ -42,16 +49,16 @@ def test_circuit_pass_through():
 
 def test_circuitbreaker_monitor():
     assert CircuitBreakerMonitor.all_closed() is True
-    assert len(list(CircuitBreakerMonitor.get_circuits())) == 5
-    assert len(list(CircuitBreakerMonitor.get_closed())) == 5
+    assert len(list(CircuitBreakerMonitor.get_circuits())) == 6
+    assert len(list(CircuitBreakerMonitor.get_closed())) == 6
     assert len(list(CircuitBreakerMonitor.get_open())) == 0
 
     with raises(IOError):
         circuit_failure()
 
     assert CircuitBreakerMonitor.all_closed() is False
-    assert len(list(CircuitBreakerMonitor.get_circuits())) == 5
-    assert len(list(CircuitBreakerMonitor.get_closed())) == 4
+    assert len(list(CircuitBreakerMonitor.get_circuits())) == 6
+    assert len(list(CircuitBreakerMonitor.get_closed())) == 5
     assert len(list(CircuitBreakerMonitor.get_open())) == 1
 
 
@@ -218,3 +225,20 @@ def test_circuitbreaker_reopens_after_successful_calls(mock_remote):
     assert circuit_threshold_2_timeout_1()
     assert circuit_threshold_2_timeout_1()
     assert circuit_threshold_2_timeout_1()
+
+
+@patch("test_functional.pseudo_remote_call", return_value=True)
+def test_circuitbreaker_handles_generator_functions(mock_remote):
+    # type: (Mock) -> None
+    circuitbreaker = CircuitBreakerMonitor.get("circuit_generator_failure")
+    assert circuitbreaker.closed
+
+    with raises(IOError):
+        list(circuit_generator_failure())
+
+    assert circuitbreaker.opened
+
+    with raises(CircuitBreakerError):
+        list(circuit_generator_failure())
+
+    mock_remote.assert_called_once()

tests/test_unit.py

@@ -47,23 +47,27 @@ def test_circuitbreaker_should_clear_last_exception_on_success_call():
 def test_circuitbreaker_should_call_fallback_function_if_open():
     fallback = Mock(return_value=True)
 
-    func = Mock(return_value=False)
+    func = Mock(return_value=False, __name__="Mock") # attribute __name__ required for 2.7 compat with functools.wraps
 
     CircuitBreaker.opened = lambda self: True
 
     cb = CircuitBreaker(name='WithFallback', fallback_function=fallback)
-    cb.call(func)
+    decorated_func = cb.decorate(func)
+
+    decorated_func()
     fallback.assert_called_once_with()
 
 def test_circuitbreaker_should_not_call_function_if_open():
     fallback = Mock(return_value=True)
 
-    func = Mock(return_value=False)
+    func = Mock(return_value=False, __name__="Mock") # attribute __name__ required for 2.7 compat with functools.wraps
 
     CircuitBreaker.opened = lambda self: True
 
     cb = CircuitBreaker(name='WithFallback', fallback_function=fallback)
-    assert cb.call(func) == fallback.return_value
+    decorated_func = cb.decorate(func)
+
+    assert decorated_func() == fallback.return_value
     assert not func.called