Skip to content

Commit

Permalink
tests: make Std{err,out}Buffer context managers stackable
Browse files Browse the repository at this point in the history
Allow use StdoutBuffer context manager inside another similar context
manager. This allows running the tests by a test runner which collect
the output using similar trick.
  • Loading branch information
marmarek committed Apr 30, 2017
1 parent 54dcec2 commit cdb9205
Showing 1 changed file with 6 additions and 2 deletions.
8 changes: 6 additions & 2 deletions qubesmgmt/tests/tools/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,31 +25,35 @@

class StdoutBuffer(object):
def __init__(self):
self.orig_stdout = None
if sys.version_info[0] >= 3:
self.stdout = io.StringIO()
else:
self.stdout = io.BytesIO()

def __enter__(self):
self.orig_stdout = sys.stdout
sys.stdout = self.stdout
return self.stdout

def __exit__(self, exc_type, exc_val, exc_tb):
sys.stdout = sys.__stdout__
sys.stdout = self.orig_stdout
return False


class StderrBuffer(object):
def __init__(self):
self.orig_stderr = None
if sys.version_info[0] >= 3:
self.stderr = io.StringIO()
else:
self.stderr = io.BytesIO()

def __enter__(self):
self.orig_stderr = sys.stderr
sys.stderr = self.stderr
return self.stderr

def __exit__(self, exc_type, exc_val, exc_tb):
sys.stderr = sys.__stderr__
sys.stderr = self.orig_stderr
return False

0 comments on commit cdb9205

Please sign in to comment.