Merge "Make consumer_for a context manager"

marconi/common/pipeline.py

@@ -30,6 +30,8 @@ At least one of the stages has to implement the calling method. If none of
 them do, an AttributeError exception will be raised.
 """
 
+import contextlib
+
 import six
 
 from marconi.common import decorators
@@ -49,8 +51,10 @@ class Pipeline(object):
 
     @decorators.cached_getattr
     def __getattr__(self, name):
-        return self.consumer_for(name)
+        with self.consumer_for(name) as consumer:
+            return consumer
 
+    @contextlib.contextmanager
     def consumer_for(self, method):
         """Creates a closure for `method`
 
@@ -107,4 +111,4 @@ class Pipeline(object):
                 LOG.error(msg)
                 raise AttributeError(msg)
 
-        return consumer
+        yield consumer

tests/unit/common/test_pipeline.py

@@ -63,7 +63,7 @@ class TestPipeLine(base.TestBase):
                                            SecondClass()])
 
     def test_attribute_error(self):
-        consumer = self.pipeline.consumer_for('does_not_exist')
+        consumer = self.pipeline.does_not_exist
         self.assertRaises(AttributeError, consumer)
 
     def test_with_args(self):
@@ -87,3 +87,9 @@ class TestPipeLine(base.TestBase):
 
     def test_calls_the_latest(self):
         self.assertTrue(self.pipeline.calls_the_latest())
+
+    def test_pipeline_context_manager(self):
+        ctxt = self.pipeline.consumer_for('does_nothing')
+
+        with ctxt as consumer:
+            self.assertIsNone(consumer())