Refactor capturing stderr into a context manager

2 files changed, 9 insertions, 19 deletions

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):