test: Follow-up upcoming caucase API changes.

kedifa/test.py

@@ -61,17 +61,6 @@ def findFreeTCPPort(ip=''):
     return str(s.getsockname()[1])
 
 
-@contextlib.contextmanager
-def captured_output():
-  new_out, new_err = StringIO(), StringIO()
-  old_out, old_err = sys.stdout, sys.stderr
-  try:
-    sys.stdout, sys.stderr = new_out, new_err
-    yield sys.stdout, sys.stderr
-  finally:
-    sys.stdout, sys.stderr = old_out, old_err
-
-
 class KedifaMixin(object):
   def setUp(self):
     self.testdir = tempfile.mkdtemp()
@@ -200,10 +189,15 @@ class KedifaMixinCaucase(KedifaMixin):
     self.cas = cas.split()
 
     kedifa_key_pem, csr_file = self.generateCSR(unicode(common_name))
-    with captured_output() as (out, err):
-      caucase.cli.main(argv=self.cas + [
+    out = StringIO()
+    err = StringIO()
+    caucase.cli.main(
+      argv=self.cas + [
         '--send-csr', csr_file
-      ])
+      ],
+      stdout=out,
+      stderr=err,
+    )
 
     self.assertEqual('', err.getvalue().strip())
     output = out.getvalue().strip()
@@ -439,10 +433,15 @@ class KedifaIntegrationTest(KedifaMixinCaucase, unittest.TestCase):
       raise
 
   def revokeCaucaseServiceCertifice(self):
-    with captured_output() as (out, err):
-      caucase.cli.main(argv=self.cas + [
+    out = StringIO()
+    err = StringIO()
+    caucase.cli.main(
+      argv=self.cas + [
         '--revoke-crt', self.client_key_pem, self.client_key_pem
-      ])
+      ],
+      stdout=out,
+      stderr=err,
+    )
 
     self.assertEqual('', out.getvalue().strip())
     self.assertEqual('', err.getvalue().strip())