diff options
Diffstat (limited to 'test/units/mock/procenv.py')
| -rw-r--r-- | test/units/mock/procenv.py | 31 |
1 files changed, 20 insertions, 11 deletions
diff --git a/test/units/mock/procenv.py b/test/units/mock/procenv.py index e9d470c079..6cf69a7acc 100644 --- a/test/units/mock/procenv.py +++ b/test/units/mock/procenv.py @@ -36,18 +36,22 @@ def swap_stdin_and_argv(stdin_data='', argv_data=tuple()): context manager that temporarily masks the test runner's values for stdin and argv """ real_stdin = sys.stdin + real_argv = sys.argv if PY3: - sys.stdin = StringIO(stdin_data) - sys.stdin.buffer = BytesIO(to_bytes(stdin_data)) + fake_stream = StringIO(stdin_data) + fake_stream.buffer = BytesIO(to_bytes(stdin_data)) else: - sys.stdin = BytesIO(to_bytes(stdin_data)) + fake_stream = BytesIO(to_bytes(stdin_data)) - real_argv = sys.argv - sys.argv = argv_data - yield - sys.stdin = real_stdin - sys.argv = real_argv + try: + sys.stdin = fake_stream + sys.argv = argv_data + + yield + finally: + sys.stdin = real_stdin + sys.argv = real_argv @contextmanager @@ -56,13 +60,18 @@ def swap_stdout(): context manager that temporarily replaces stdout for tests that need to verify output """ old_stdout = sys.stdout + if PY3: fake_stream = StringIO() else: fake_stream = BytesIO() - sys.stdout = fake_stream - yield fake_stream - sys.stdout = old_stdout + + try: + sys.stdout = fake_stream + + yield fake_stream + finally: + sys.stdout = old_stdout class ModuleTestCase(unittest.TestCase): |