move annotate_failure to test.__init__

make this useful helper available as a function for all tests to use

Change-Id: I428c6fe899aa3877f2575aad224e275060174535

test/__init__.py

@@ -17,7 +17,11 @@
 # The code below enables nosetests to work with i18n _() blocks
 from __future__ import print_function
 import sys
+from contextlib import contextmanager
+
 import os
+from six import reraise
+
 try:
     from unittest.util import safe_repr
 except ImportError:
@@ -86,3 +90,26 @@ def listen_zero():
     sock.bind(("127.0.0.1", 0))
     sock.listen(50)
     return sock
+
+
+@contextmanager
+def annotate_failure(msg):
+    """
+    Catch AssertionError and annotate it with a message. Useful when making
+    assertions in a loop where the message can indicate the loop index or
+    richer context about the failure.
+
+    :param msg: A message to be prefixed to the AssertionError message.
+    """
+    try:
+        yield
+    except AssertionError as err:
+        err_typ, err_val, err_tb = sys.exc_info()
+        if err_val.args:
+            msg = '%s Failed with %s' % (msg, err_val.args[0])
+            err_val.args = (msg, ) + err_val.args[1:]
+        else:
+            # workaround for some IDE's raising custom AssertionErrors
+            err_val = '%s Failed with %s' % (msg, err)
+            err_typ = AssertionError
+        reraise(err_typ, err_val, err_tb)

test/probe/common.py

@@ -15,8 +15,6 @@
 
 from __future__ import print_function
 
-from contextlib import contextmanager
-
 import os
 from subprocess import Popen, PIPE
 import sys
@@ -483,16 +481,6 @@ class ProbeTest(unittest.TestCase):
         finally:
             shutil.rmtree(tempdir)
 
-    @contextmanager
-    def annotate_failure(self, msg):
-        try:
-            yield
-        except AssertionError:
-            err_typ, err_val, err_tb = sys.exc_info()
-            msg = '%s Failed with %s' % (msg, err_val.args[0])
-            err_val.args = (msg, ) + err_val.args[1:]
-            raise err_typ, err_val, err_tb
-
 
 class ReplProbeTest(ProbeTest):
 

test/probe/test_sharder.py

@@ -28,6 +28,7 @@ from swift.common import utils
 from swift.common.manager import Manager
 from swiftclient import client, get_auth, ClientException
 
+from test import annotate_failure
 from test.probe.brain import BrainSplitter
 from test.probe.common import ReplProbeTest, get_server_number
 
@@ -764,7 +765,7 @@ class TestContainerSharding(ReplProbeTest):
             object_counts = []
             bytes_used = []
             for node_id, node_data in node_data.items():
-                with self.annotate_failure('Node id %s.' % node_id):
+                with annotate_failure('Node id %s.' % node_id):
                     check_node_data(
                         node_data, exp_shard_hdrs, exp_obj_count,
                         expected_shards)
@@ -797,7 +798,7 @@ class TestContainerSharding(ReplProbeTest):
         root_nodes_data = self.direct_get_container_shard_ranges()
         self.assertEqual(3, len(root_nodes_data))
         for node_id, node_data in root_nodes_data.items():
-            with self.annotate_failure('Node id %s.' % node_id):
+            with annotate_failure('Node id %s.' % node_id):
                 check_node_data(node_data, exp_hdrs, exp_obj_count, 2)
 
         # run updaters to update .sharded account; shard containers have not
@@ -835,7 +836,7 @@ class TestContainerSharding(ReplProbeTest):
         root_nodes_data = self.direct_get_container_shard_ranges()
         self.assertEqual(3, len(root_nodes_data))
         for node_id, node_data in root_nodes_data.items():
-            with self.annotate_failure('Node id %s.' % node_id):
+            with annotate_failure('Node id %s.' % node_id):
                 check_node_data(node_data, exp_hdrs, exp_obj_count, 2)
             range_data = node_data[1]
             self.assert_shard_range_lists_equal(
@@ -856,7 +857,7 @@ class TestContainerSharding(ReplProbeTest):
         self.assertEqual(3, len(root_nodes_data))
         exp_hdrs['X-Container-Object-Count'] = str(exp_obj_count)
         for node_id, node_data in root_nodes_data.items():
-            with self.annotate_failure('Node id %s.' % node_id):
+            with annotate_failure('Node id %s.' % node_id):
                 # NB now only *one* shard range in root
                 check_node_data(node_data, exp_hdrs, exp_obj_count, 1)