Commit fc3bf97e authored by Vincent Pelletier's avatar Vincent Pelletier

test: Use a context manager for temporary stdout overrides.

parent b4cc6236
...@@ -22,6 +22,7 @@ Test suite ...@@ -22,6 +22,7 @@ Test suite
""" """
# pylint: disable=too-many-lines, too-many-public-methods # pylint: disable=too-many-lines, too-many-public-methods
from __future__ import absolute_import from __future__ import absolute_import
import contextlib
from Cookie import SimpleCookie from Cookie import SimpleCookie
import datetime import datetime
import errno import errno
...@@ -292,6 +293,19 @@ def print_buffer_on_error(func): ...@@ -292,6 +293,19 @@ def print_buffer_on_error(func):
raise raise
return wrapper 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): class CaucaseTest(unittest.TestCase):
""" """
Test a complete caucase setup: spawn a caucase-http server on CAUCASE_NETLOC Test a complete caucase setup: spawn a caucase-http server on CAUCASE_NETLOC
...@@ -583,22 +597,19 @@ class CaucaseTest(unittest.TestCase): ...@@ -583,22 +597,19 @@ class CaucaseTest(unittest.TestCase):
Returns stdout. Returns stdout.
""" """
orig_stdout = sys.stdout with captureStdout() as stdout:
sys.stdout = stdout = BytesIO() try:
try: cli.main(
cli.main( argv=(
argv=( '--ca-url', self._caucase_url,
'--ca-url', self._caucase_url, '--ca-crt', self._client_ca_crt,
'--ca-crt', self._client_ca_crt, '--user-ca-crt', self._client_user_ca_crt,
'--user-ca-crt', self._client_user_ca_crt, '--crl', self._client_crl,
'--crl', self._client_crl, '--user-crl', self._client_user_crl,
'--user-crl', self._client_user_crl, ) + argv,
) + argv, )
) except SystemExit:
except SystemExit: pass
pass
finally:
sys.stdout = orig_stdout
return stdout.getvalue() return stdout.getvalue()
@staticmethod @staticmethod
...@@ -2380,10 +2391,8 @@ class CaucaseTest(unittest.TestCase): ...@@ -2380,10 +2391,8 @@ class CaucaseTest(unittest.TestCase):
) )
self._runClient() self._runClient()
getBytePass_orig = caucase.http.getBytePass getBytePass_orig = caucase.http.getBytePass
orig_stdout = sys.stdout
try: try:
caucase.http.getBytePass = lambda x: b'test' caucase.http.getBytePass = lambda x: b'test'
sys.stdout = stdout = BytesIO()
self.assertFalse(os.path.exists(exported_ca), exported_ca) self.assertFalse(os.path.exists(exported_ca), exported_ca)
caucase.http.manage( caucase.http.manage(
argv=( argv=(
...@@ -2394,15 +2403,16 @@ class CaucaseTest(unittest.TestCase): ...@@ -2394,15 +2403,16 @@ class CaucaseTest(unittest.TestCase):
self.assertTrue(os.path.exists(exported_ca), exported_ca) self.assertTrue(os.path.exists(exported_ca), exported_ca)
server_db2 = self._server_db + '2' server_db2 = self._server_db + '2'
self.assertFalse(os.path.exists(server_db2), server_db2) self.assertFalse(os.path.exists(server_db2), server_db2)
caucase.http.manage( with captureStdout() as stdout:
argv=( caucase.http.manage(
'--db', server_db2, argv=(
'--import-ca', exported_ca, '--db', server_db2,
'--import-crl', self._client_crl, '--import-ca', exported_ca,
# Twice, for code coverage... '--import-crl', self._client_crl,
'--import-crl', self._client_crl, # Twice, for code coverage...
), '--import-crl', self._client_crl,
) ),
)
self.assertTrue(os.path.exists(server_db2), server_db2) self.assertTrue(os.path.exists(server_db2), server_db2)
self.assertEqual( self.assertEqual(
[ [
...@@ -2412,7 +2422,6 @@ class CaucaseTest(unittest.TestCase): ...@@ -2412,7 +2422,6 @@ class CaucaseTest(unittest.TestCase):
stdout.getvalue().splitlines(), stdout.getvalue().splitlines(),
) )
finally: finally:
sys.stdout = orig_stdout
caucase.http.getBytePass = getBytePass_orig caucase.http.getBytePass = getBytePass_orig
def testWSGIBase(self): def testWSGIBase(self):
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment