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,22 +597,19 @@ class CaucaseTest(unittest.TestCase):
 
     Returns stdout.
     """
-    orig_stdout = sys.stdout
-    sys.stdout = stdout = BytesIO()
-    try:
-      cli.main(
-        argv=(
-          '--ca-url', self._caucase_url,
-          '--ca-crt', self._client_ca_crt,
-          '--user-ca-crt', self._client_user_ca_crt,
-          '--crl', self._client_crl,
-          '--user-crl', self._client_user_crl,
-        ) + argv,
-      )
-    except SystemExit:
-      pass
-    finally:
-      sys.stdout = orig_stdout
+    with captureStdout() as stdout:
+      try:
+        cli.main(
+          argv=(
+            '--ca-url', self._caucase_url,
+            '--ca-crt', self._client_ca_crt,
+            '--user-ca-crt', self._client_user_ca_crt,
+            '--crl', self._client_crl,
+            '--user-crl', self._client_user_crl,
+          ) + argv,
+        )
+      except SystemExit:
+        pass
     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,15 +2403,16 @@ 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)
-      caucase.http.manage(
-        argv=(
-          '--db', server_db2,
-          '--import-ca', exported_ca,
-          '--import-crl', self._client_crl,
-          # Twice, for code coverage...
-          '--import-crl', self._client_crl,
-        ),
-      )
+      with captureStdout() as stdout:
+        caucase.http.manage(
+          argv=(
+            '--db', server_db2,
+            '--import-ca', exported_ca,
+            '--import-crl', self._client_crl,
+            # Twice, for code coverage...
+            '--import-crl', self._client_crl,
+          ),
+        )
       self.assertTrue(os.path.exists(server_db2), server_db2)
       self.assertEqual(
         [
@@ -2412,7 +2422,6 @@ class CaucaseTest(unittest.TestCase):
         stdout.getvalue().splitlines(),
       )
     finally:
-      sys.stdout = orig_stdout
       caucase.http.getBytePass = getBytePass_orig
 
   def testWSGIBase(self):