trap-failure processor

nkvoll · Sep 16, 2011 · 52f679a · 52f679a
1 parent 07a17b7
commit 52f679a
Show file tree

Hide file tree

Showing 3 changed files with 89 additions and 4 deletions.
diff --git a/doc/reference/processors.rst b/doc/reference/processors.rst
@@ -202,6 +202,13 @@ stop
 .. autoclass:: Stopper
 
 
+.. _trap-failure:
+
+trap-failure
+^^^^^^^^^^^^
+.. autoclass:: TrapFailure
+
+
 .. _wait:
 
 wait

diff --git a/piped/processors/test/test_util_processors.py b/piped/processors/test/test_util_processors.py
@@ -13,6 +13,10 @@
 from piped.processors import util_processors
 
 
+class FakeError(exceptions.PipedError):
+    pass
+
+
 class TestMergeWithDict(unittest.TestCase):
 
     def test_merge_with_dict(self):
@@ -217,6 +221,51 @@ def test_baton_cleaner_err_when_misconfigured(self):
             pass
 
 
+class TestTrapFailure(unittest.TestCase):
+
+    def _create_processor(self, **kwargs):
+        return util_processors.TrapFailure(**kwargs)
+
+    def test_simple_trapping(self):
+        processor = self._create_processor(error_types=reflect.fullyQualifiedName(FakeError), output_path='trapped')
+
+        try:
+            raise FakeError('test')
+        except FakeError as fe:
+            baton = processor.process(dict())
+            self.assertEquals(baton['trapped'], FakeError)
+
+    def test_not_trapping_unexpected_exceptions(self):
+        processor = self._create_processor(error_types=reflect.fullyQualifiedName(FakeError))
+
+        try:
+            raise exceptions.PipedError('test')
+        except exceptions.PipedError as pe:
+            try:
+                processor.process(None)
+            except failure.Failure as reason:
+                self.assertEquals(reason.type, exceptions.PipedError)
+                self.assertEquals(reason.value, pe)
+            else:
+                self.fail('Expected a failure to be raised.')
+
+    def test_no_current_failure(self):
+        processor = self._create_processor(error_types=reflect.fullyQualifiedName(FakeError))
+
+        self.assertRaises(failure.NoCurrentExceptionError, processor.process, None)
+
+    def test_trapping_multiple_types(self):
+        error_types = [reflect.fullyQualifiedName(FakeError), reflect.fullyQualifiedName(exceptions.ConfigurationError)]
+        processor = self._create_processor(error_types=error_types, output_path='trapped')
+
+        for error_type in (FakeError, exceptions.ConfigurationError):
+            try:
+                raise error_type('test')
+            except error_type as fe:
+                baton = processor.process(dict())
+                self.assertEquals(baton['trapped'], error_type)
+
+
 class TestFlattenDictionaryList(unittest.TestCase):
 
     @defer.inlineCallbacks
@@ -586,10 +635,6 @@ def process(self, baton):
         defer.returnValue([baton])
 
 
-class FakeError(exceptions.PipedError):
-    pass
-
-
 class FakeFailingPipeline(FakePipeline):
 
     def process(self, baton):

diff --git a/piped/processors/util_processors.py b/piped/processors/util_processors.py
@@ -278,6 +278,39 @@ def process(self, baton):
         f.printTraceback()
         return baton
 
+class TrapFailure(base.Processor):
+    """ Traps failures of the specified types.
+
+    If the encountered exception is not one of the expected exception types, this
+    processor will raise the original exception, preserving the traceback.
+    """
+    name = 'trap-failure'
+    interface.classProvides(processing.IProcessor)
+
+    def __init__(self, error_types, output_path=None, *a, **kw):
+        """
+        :param error_types: A single or a list of fully qualified exception class
+            names that should be trapped.
+        :param output_path: If one of the expected error types are trapped, this
+            value will be set to the matching error type.
+        """
+        super(TrapFailure, self).__init__(*a, **kw)
+
+        if not isinstance(error_types, (list, tuple)):
+            error_types = [error_types]
+
+        for i, error_type in enumerate(error_types):
+            error_types[i] = reflect.namedAny(error_type)
+
+        self.error_types = error_types
+        self.output_path = output_path
+
+    def process(self, baton):
+        f = failure.Failure()
+        trapped = f.trap(*self.error_types)
+        baton = self.get_resulting_baton(baton, self.output_path, trapped)
+        return baton
+
 
 class FlattenDictionaryList(base.InputOutputProcessor):
     """ Reduce a list of dictionaries to a list of values, given a key