util.failover: init

util/failover.py

@@ -0,0 +1,46 @@
+import logging
+
+from functools import wraps
+
+
+logger = logging.getLogger(__name__)
+
+
+class FailoverException(Exception):
+  """ Exception raised when an operation should be retried by the failover decorator. """
+  def __init__(self, message):
+    super(FailoverException, self).__init__()
+    self.message = message
+
+def failover(func):
+  """ Wraps a function such that it can be retried on specified failures.
+      Raises FailoverException when all failovers are exhausted.
+      Example:
+
+      @failover
+      def get_google(scheme, use_www=False):
+        www = 'www.' if use_www else ''
+        r = requests.get(scheme + '://' + www + 'google.com')
+        if r.status_code != 200:
+          raise FailoverException('non 200 response from Google' )
+        return r
+
+      def GooglePingTest():
+        r = get_google(
+          (('http'), {'use_www': False}),
+          (('http'), {'use_www': True}),
+          (('https'), {'use_www': False}),
+          (('https'), {'use_www': True}),
+        )
+        print('Successfully contacted ' + r.url)
+  """
+  @wraps(func)
+  def wrapper(*args_sets):
+    for arg_set in args_sets:
+      try:
+        return func(*arg_set[0], **arg_set[1])
+      except FailoverException as ex:
+        logger.debug('failing over: %s', ex.message)
+        continue
+    raise FailoverException('exhausted all possible failovers')
+  return wrapper

util/test/test_failover.py

@@ -0,0 +1,44 @@
+import pytest
+
+from util.failover import failover, FailoverException
+
+
+class Counter(object):
+  """ Wraps a counter in an object so that it'll be passed by reference. """
+  def __init__(self):
+    self.calls = 0
+
+  def increment(self):
+    self.calls += 1
+
+
+@failover
+def my_failover_func(i, should_raise=None):
+  """ Increments a counter and raises an exception when told. """
+  i.increment()
+  if should_raise is not None:
+    raise should_raise()
+  raise FailoverException('incrementing')
+
+
+@pytest.mark.parametrize('stop_on,exception', [
+  (10, None),
+  (5, IndexError),
+])
+def test_readonly_failover(stop_on, exception):
+  """ Generates failover arguments and checks against a counter to ensure that
+      the failover function has been called the proper amount of times and stops
+      at unhandled exceptions.
+  """
+  counter = Counter()
+  arg_sets = []
+  for i in xrange(stop_on):
+    should_raise = exception if exception is not None and i == stop_on-1 else None
+    arg_sets.append(((counter,), {'should_raise': should_raise}))
+
+  if exception is not None:
+    with pytest.raises(exception):
+      my_failover_func(*arg_sets)
+  else:
+    my_failover_func(*arg_sets)
+  assert counter.calls == stop_on