test: Use a context manager for temporary stdout overrides.

caucase/test.py

@@ -22,6 +22,7 @@ Test suite
 """
 # pylint: disable=too-many-lines, too-many-public-methods
 from __future__ import absolute_import
+import contextlib
 from Cookie import SimpleCookie
 import datetime
 import errno
@@ -292,6 +293,19 @@ def print_buffer_on_error(func):
       raise
   return wrapper
 
+@contextlib.contextmanager
+def captureStdout():
+  """
+  Replace stdout with a BytesIO object for the duration of the context manager,
+  and provide it to caller.
+  """
+  orig_stdout = sys.stdout
+  sys.stdout = stdout = BytesIO()
+  try:
+    yield stdout
+  finally:
+    sys.stdout = orig_stdout
+
 class CaucaseTest(unittest.TestCase):
   """
   Test a complete caucase setup: spawn a caucase-http server on CAUCASE_NETLOC
@@ -583,8 +597,7 @@ class CaucaseTest(unittest.TestCase):
 
     Returns stdout.
     """
-    orig_stdout = sys.stdout
-    sys.stdout = stdout = BytesIO()
+    with captureStdout() as stdout:
       try:
         cli.main(
           argv=(
@@ -597,8 +610,6 @@ class CaucaseTest(unittest.TestCase):
         )
       except SystemExit:
         pass
-    finally:
-      sys.stdout = orig_stdout
     return stdout.getvalue()
 
   @staticmethod
@@ -2380,10 +2391,8 @@ class CaucaseTest(unittest.TestCase):
     )
     self._runClient()
     getBytePass_orig = caucase.http.getBytePass
-    orig_stdout = sys.stdout
     try:
       caucase.http.getBytePass = lambda x: b'test'
-      sys.stdout = stdout = BytesIO()
       self.assertFalse(os.path.exists(exported_ca), exported_ca)
       caucase.http.manage(
         argv=(
@@ -2394,6 +2403,7 @@ class CaucaseTest(unittest.TestCase):
       self.assertTrue(os.path.exists(exported_ca), exported_ca)
       server_db2 = self._server_db + '2'
       self.assertFalse(os.path.exists(server_db2), server_db2)
+      with captureStdout() as stdout:
         caucase.http.manage(
           argv=(
             '--db', server_db2,
@@ -2412,7 +2422,6 @@ class CaucaseTest(unittest.TestCase):
         stdout.getvalue().splitlines(),
       )
     finally:
-      sys.stdout = orig_stdout
       caucase.http.getBytePass = getBytePass_orig
 
   def testWSGIBase(self):