1# (c) 2016, Matt Davis <mdavis@ansible.com>
2# (c) 2016, Toshio Kuratomi <tkuratomi@ansible.com>
3#
4# This file is part of Ansible
5#
6# Ansible is free software: you can redistribute it and/or modify
7# it under the terms of the GNU General Public License as published by
8# the Free Software Foundation, either version 3 of the License, or
9# (at your option) any later version.
10#
11# Ansible is distributed in the hope that it will be useful,
12# but WITHOUT ANY WARRANTY; without even the implied warranty of
13# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14# GNU General Public License for more details.
15#
16# You should have received a copy of the GNU General Public License
17# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.
18
19# Make coding more python3-ish
20from __future__ import (absolute_import, division, print_function)
21__metaclass__ = type
22
23import sys
24import json
25
26from contextlib import contextmanager
27from io import BytesIO, StringIO
28from ansible_collections.ansible.libvirt.tests.unit.compat import unittest
29from ansible.module_utils.six import PY3
30from ansible.module_utils._text import to_bytes
31
32
33@contextmanager
34def swap_stdin_and_argv(stdin_data='', argv_data=tuple()):
35    """
36    context manager that temporarily masks the test runner's values for stdin and argv
37    """
38    real_stdin = sys.stdin
39    real_argv = sys.argv
40
41    if PY3:
42        fake_stream = StringIO(stdin_data)
43        fake_stream.buffer = BytesIO(to_bytes(stdin_data))
44    else:
45        fake_stream = BytesIO(to_bytes(stdin_data))
46
47    try:
48        sys.stdin = fake_stream
49        sys.argv = argv_data
50
51        yield
52    finally:
53        sys.stdin = real_stdin
54        sys.argv = real_argv
55
56
57@contextmanager
58def swap_stdout():
59    """
60    context manager that temporarily replaces stdout for tests that need to verify output
61    """
62    old_stdout = sys.stdout
63
64    if PY3:
65        fake_stream = StringIO()
66    else:
67        fake_stream = BytesIO()
68
69    try:
70        sys.stdout = fake_stream
71
72        yield fake_stream
73    finally:
74        sys.stdout = old_stdout
75
76
77class ModuleTestCase(unittest.TestCase):
78    def setUp(self, module_args=None):
79        if module_args is None:
80            module_args = {'_ansible_remote_tmp': '/tmp', '_ansible_keep_remote_files': False}
81
82        args = json.dumps(dict(ANSIBLE_MODULE_ARGS=module_args))
83
84        # unittest doesn't have a clean place to use a context manager, so we have to enter/exit manually
85        self.stdin_swap = swap_stdin_and_argv(stdin_data=args)
86        self.stdin_swap.__enter__()
87
88    def tearDown(self):
89        # unittest doesn't have a clean place to use a context manager, so we have to enter/exit manually
90        self.stdin_swap.__exit__(None, None, None)
91