Attempt retry hook for flaky regression test cases.

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)

regression/datastore.py

@@ -24,6 +24,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'
@@ -33,6 +34,8 @@
 
 class TestDatastore(unittest2.TestCase):
 
+    __metaclass__ = RetryTestsMetaclass
+
     def setUp(self):
         self.case_entities_to_delete = []
 

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,37 @@
 """
 
 
+class RetryTestsMetaclass(type):
+
+    NUM_RETRIES = 2
+
+    @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:
+                    num_attempts += 1
+                    if num_attempts == RetryTestsMetaclass.NUM_RETRIES:
+                        raise
+
+        return retry_function
+
+    def __new__(cls, name, bases, attrs):
+        new_attrs = {}
+        for attr_name, value in attrs.items():
+            new_attrs[attr_name] = cls._wrap_class_attr(value)
+
+        return super(RetryTestsMetaclass, cls).__new__(
+            cls, 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()