Attempt retry hook for flaky regression test cases.

dhermes · dhermes · commit 012fd2c43ea2 · 2015-01-15T13:50:41.000-08:00
Fixes googleapis#531. To "test" that this works, feel free to add a test case like: x = 0 def test_retry(self): # Feel free to vary 3 higher and higher, should always be # NUM_RETRIES in the final error message. if self.x < 3: self.x += 1 self.assertEqual(self.x, object()) # Fails else: self.assertTrue(True)
diff --git a/regression/datastore.py b/regression/datastore.py
@@ -20,6 +20,7 @@
 # This assumes the command is being run via tox hence the
 # repository root is the current directory.
 from regression import populate_datastore
+from regression.regression_utils import RetryTestsMetaclass
 
 
 datastore._DATASET_ENV_VAR_NAME = 'GCLOUD_TESTS_DATASET_ID'
@@ -28,6 +29,8 @@
 
 class TestDatastore(unittest2.TestCase):
 
+    __metaclass__ = RetryTestsMetaclass
+
     def setUp(self):
         self.case_entities_to_delete = []
 
diff --git a/regression/regression_utils.py b/regression/regression_utils.py
@@ -15,6 +15,7 @@
 from __future__ import print_function
 import os
 import sys
+import types
 
 from gcloud import storage
 
@@ -31,6 +32,38 @@
 """
 
 
+class RetryTestsMetaclass(type):
+
+    NUM_RETRIES = 2
+    FLAKY_ERROR_CLASSES = (AssertionError,)
+
+    @staticmethod
+    def _wrap_class_attr(class_attr):
+        if not (isinstance(class_attr, types.FunctionType) and
+                class_attr.__name__.startswith('test_')):
+            return class_attr
+
+        def retry_function(self):
+            num_attempts = 0
+            while num_attempts < RetryTestsMetaclass.NUM_RETRIES:
+                try:
+                    return class_attr(self)
+                except RetryTestsMetaclass.FLAKY_ERROR_CLASSES:
+                    num_attempts += 1
+                    if num_attempts == RetryTestsMetaclass.NUM_RETRIES:
+                        raise
+
+        return retry_function
+
+    def __new__(mcs, name, bases, attrs):
+        new_attrs = {}
+        for attr_name, value in attrs.items():
+            new_attrs[attr_name] = mcs._wrap_class_attr(value)
+
+        return super(RetryTestsMetaclass, mcs).__new__(
+            mcs, name, bases, new_attrs)
+
+
 def get_environ(require_datastore=False, require_storage=False):
     if require_datastore:
         if DATASET_ID is None or not os.path.isfile(CREDENTIALS):
diff --git a/regression/storage.py b/regression/storage.py
@@ -47,6 +47,8 @@ def tearDownModule():
 
 class TestStorage(unittest2.TestCase):
 
+    __metaclass__ = regression_utils.RetryTestsMetaclass
+
     @classmethod
     def setUpClass(cls):
         cls.connection = regression_utils.get_storage_connection()
