mirror of
https://github.com/thegeeklab/ansible-later.git
synced 2024-11-30 00:30:35 +00:00
75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
|
import sys
|
||
|
|
|
||
|
|
from testfixtures.comparison import compare
|
||
|
|
from testfixtures.compat import StringIO
|
||
|
|
|
||
|
|
|
||
|
|
class OutputCapture(object):
|
||
|
|
"""
|
||
|
|
A context manager for capturing output to the
|
||
|
|
:attr:`sys.stdout` and :attr:`sys.stderr` streams.
|
||
|
|
|
||
|
|
:param separate: If ``True``, ``stdout`` and ``stderr`` will be captured
|
||
|
|
separately and their expected values must be passed to
|
||
|
|
:meth:`~OutputCapture.compare`.
|
||
|
|
|
||
|
|
.. note:: If ``separate`` is passed as ``True``,
|
||
|
|
:attr:`OutputCapture.captured` will be an empty string.
|
||
|
|
"""
|
||
|
|
|
||
|
|
original_stdout = None
|
||
|
|
original_stderr = None
|
||
|
|
|
||
|
|
def __init__(self, separate=False):
|
||
|
|
self.separate = separate
|
||
|
|
|
||
|
|
def __enter__(self):
|
||
|
|
self.output = StringIO()
|
||
|
|
self.stdout = StringIO()
|
||
|
|
self.stderr = StringIO()
|
||
|
|
self.enable()
|
||
|
|
return self
|
||
|
|
|
||
|
|
def __exit__(self, *args):
|
||
|
|
self.disable()
|
||
|
|
|
||
|
|
def disable(self):
|
||
|
|
"Disable the output capture if it is enabled."
|
||
|
|
sys.stdout = self.original_stdout
|
||
|
|
sys.stderr = self.original_stderr
|
||
|
|
|
||
|
|
def enable(self):
|
||
|
|
"Enable the output capture if it is disabled."
|
||
|
|
if self.original_stdout is None:
|
||
|
|
self.original_stdout = sys.stdout
|
||
|
|
self.original_stderr = sys.stderr
|
||
|
|
if self.separate:
|
||
|
|
sys.stdout = self.stdout
|
||
|
|
sys.stderr = self.stderr
|
||
|
|
else:
|
||
|
|
sys.stdout = sys.stderr = self.output
|
||
|
|
|
||
|
|
@property
|
||
|
|
def captured(self):
|
||
|
|
"A property containing any output that has been captured so far."
|
||
|
|
return self.output.getvalue()
|
||
|
|
|
||
|
|
def compare(self, expected='', stdout='', stderr=''):
|
||
|
|
"""
|
||
|
|
Compare the captured output to that expected. If the output is
|
||
|
|
not the same, an :class:`AssertionError` will be raised.
|
||
|
|
|
||
|
|
:param expected: A string containing the expected combined output
|
||
|
|
of ``stdout`` and ``stderr``.
|
||
|
|
|
||
|
|
:param stdout: A string containing the expected output to ``stdout``.
|
||
|
|
|
||
|
|
:param stderr: A string containing the expected output to ``stderr``.
|
||
|
|
"""
|
||
|
|
for prefix, _expected, captured in (
|
||
|
|
(None, expected, self.captured),
|
||
|
|
('stdout', stdout, self.stdout.getvalue()),
|
||
|
|
('stderr', stderr, self.stderr.getvalue()),
|
||
|
|
):
|
||
|
|
compare(_expected.strip(), actual=captured.strip(), prefix=prefix)
|