From e7b08cf398acbf9560483ef5c3b4013328fbca76 Mon Sep 17 00:00:00 2001
From: Jakub Stasiak <jakub@stasiak.at>
Date: Wed, 15 Oct 2014 23:58:18 +0100
Subject: [PATCH] Refactor capturing stderr into a context manager

---
 eventlet/support/__init__.py | 14 +++++++++++++-
 tests/debug_test.py          | 13 ++++---------
 tests/wsgi_test.py           | 15 +++++----------
 3 files changed, 22 insertions(+), 20 deletions(-)

diff --git a/eventlet/support/__init__.py b/eventlet/support/__init__.py
index 73b95d3..4c2b75d 100644
--- a/eventlet/support/__init__.py
+++ b/eventlet/support/__init__.py
@@ -1,6 +1,7 @@
 import sys
+from contextlib import contextmanager
 
-from eventlet.support import greenlets
+from eventlet.support import greenlets, six
 
 
 def get_errno(exc):
@@ -41,3 +42,14 @@ else:
         return b.decode(encoding)
 
 PY33 = sys.version_info[:2] == (3, 3)
+
+@contextmanager
+def capture_stderr():
+    stream = six.StringIO()
+    original = sys.stderr
+    try:
+        sys.stderr = stream
+        yield stream
+    finally:
+        sys.stderr = original
+        stream.seek(0)
diff --git a/tests/debug_test.py b/tests/debug_test.py
index 67617f4..8710dc0 100644
--- a/tests/debug_test.py
+++ b/tests/debug_test.py
@@ -2,7 +2,7 @@ import sys
 from unittest import TestCase
 
 from eventlet import debug
-from eventlet.support import six
+from eventlet.support import capture_stderr, six
 from tests import LimitedTestCase, main
 import eventlet
 
@@ -107,10 +107,7 @@ class TestDebug(LimitedTestCase):
             s.recv(1)
             {}[1]  # keyerror
 
-        fake = six.StringIO()
-        orig = sys.stderr
-        sys.stderr = fake
-        try:
+        with capture_stderr() as fake:
             gt = eventlet.spawn(hurl, client_2)
             eventlet.sleep(0)
             client.send(b' ')
@@ -118,10 +115,8 @@ class TestDebug(LimitedTestCase):
             # allow the "hurl" greenlet to trigger the KeyError
             # not sure why the extra context switch is needed
             eventlet.sleep(0)
-        finally:
-            sys.stderr = orig
-            self.assertRaises(KeyError, gt.wait)
-            debug.hub_exceptions(False)
+        self.assertRaises(KeyError, gt.wait)
+        debug.hub_exceptions(False)
         # look for the KeyError exception in the traceback
         assert 'KeyError: 1' in fake.getvalue(), "Traceback not in:\n" + fake.getvalue()
 
diff --git a/tests/wsgi_test.py b/tests/wsgi_test.py
index cdb576a..c31d79c 100644
--- a/tests/wsgi_test.py
+++ b/tests/wsgi_test.py
@@ -17,7 +17,7 @@ from eventlet.green import subprocess
 from eventlet import greenio
 from eventlet import greenthread
 from eventlet import support
-from eventlet.support import bytes_to_str, six
+from eventlet.support import bytes_to_str, capture_stderr, six
 from eventlet import tpool
 from eventlet import wsgi
 
@@ -835,11 +835,8 @@ class TestHttpd(_TestBase):
         listener = greensocket.socket()
         listener.bind(('localhost', 0))
         # NOT calling listen, to trigger the error
-        self.logfile = six.StringIO()
-        self.spawn_server(sock=listener)
-        old_stderr = sys.stderr
-        try:
-            sys.stderr = self.logfile
+        with capture_stderr() as log:
+            self.spawn_server(sock=listener)
             eventlet.sleep(0)  # need to enter server loop
             try:
                 eventlet.connect(('localhost', self.port))
@@ -847,10 +844,8 @@ class TestHttpd(_TestBase):
             except socket.error as exc:
                 self.assertEqual(support.get_errno(exc), errno.ECONNREFUSED)
 
-            log_content = self.logfile.getvalue()
-            assert 'Invalid argument' in log_content, log_content
-        finally:
-            sys.stderr = old_stderr
+        log_content = log.getvalue()
+        assert 'Invalid argument' in log_content, log_content
         debug.hub_exceptions(False)
 
     def test_026_log_format(self):