Attempt retry hook for flaky regression test cases.

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 = []
 
@@ -77,7 +80,8 @@ def _get_post(self, id_or_name=None, post_content=None):
 
     def _generic_test_post(self, name=None, key_id=None):
         entity = self._get_post(id_or_name=(name or key_id))
-        datastore.put([entity])
+        with datastore.Transaction():
+            datastore.put([entity])
 
         # Register entity to be deleted.
         self.case_entities_to_delete.append(entity)
@@ -146,7 +150,8 @@ def test_save_key_self_reference(self):
         entity['fullName'] = u'Full name'
         entity['linkedTo'] = key  # Self reference.
 
-        datastore.put([entity])
+        with datastore.Transaction():
+            datastore.put([entity])
         self.case_entities_to_delete.append(entity)
 
         query = datastore.Query(kind='Person')

regression/regression_utils.py

@@ -15,8 +15,10 @@
 from __future__ import print_function
 import os
 import sys
+import types
 
 from gcloud import storage
+from gcloud.storage.exceptions import NotFound
 
 
 # From shell environ. May be None.
@@ -31,6 +33,38 @@
 """
 
 
+class RetryTestsMetaclass(type):
+
+    NUM_RETRIES = 2
+    FLAKY_ERROR_CLASSES = (NotFound,)
+
+    @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):

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()