From 8fa010db9890cb21559e7a01e23bc9b725f5ccce Mon Sep 17 00:00:00 2001
From: Flavio Percoco <flaper87@gmail.com>
Date: Fri, 29 Nov 2013 21:30:49 +0100
Subject: [PATCH] Make consumer_for a context manager

This patch converts consumer_for into a `contextmanager` decorated
method and adapts `__getattribute__` to support this change. A new
method was not added to avoid clashes with method names in the stages.

Change-Id: I6c0727e86d9ee8d6bd34c312d180040c8baf34ba
---
 marconi/common/pipeline.py         | 8 ++++++--
 tests/unit/common/test_pipeline.py | 8 +++++++-
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/marconi/common/pipeline.py b/marconi/common/pipeline.py
index 48ae8f93c..7a446d647 100644
--- a/marconi/common/pipeline.py
+++ b/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
@@ -48,8 +50,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`
 
@@ -106,4 +110,4 @@ class Pipeline(object):
                 LOG.error(msg)
                 raise AttributeError(msg)
 
-        return consumer
+        yield consumer
diff --git a/tests/unit/common/test_pipeline.py b/tests/unit/common/test_pipeline.py
index 9abff7003..29d6c733f 100644
--- a/tests/unit/common/test_pipeline.py
+++ b/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())