1import atexit
2import faulthandler
3import os
4import signal
5import sys
6import unittest
7from test import support
8from test.support.os_helper import TESTFN_UNDECODABLE, FS_NONASCII
9try:
10    import gc
11except ImportError:
12    gc = None
13
14from test.libregrtest.utils import (setup_unraisable_hook,
15                                    setup_threading_excepthook)
16
17
18UNICODE_GUARD_ENV = "PYTHONREGRTEST_UNICODE_GUARD"
19
20
21def setup_tests(ns):
22    try:
23        stderr_fd = sys.__stderr__.fileno()
24    except (ValueError, AttributeError):
25        # Catch ValueError to catch io.UnsupportedOperation on TextIOBase
26        # and ValueError on a closed stream.
27        #
28        # Catch AttributeError for stderr being None.
29        stderr_fd = None
30    else:
31        # Display the Python traceback on fatal errors (e.g. segfault)
32        faulthandler.enable(all_threads=True, file=stderr_fd)
33
34        # Display the Python traceback on SIGALRM or SIGUSR1 signal
35        signals = []
36        if hasattr(signal, 'SIGALRM'):
37            signals.append(signal.SIGALRM)
38        if hasattr(signal, 'SIGUSR1'):
39            signals.append(signal.SIGUSR1)
40        for signum in signals:
41            faulthandler.register(signum, chain=True, file=stderr_fd)
42
43    _adjust_resource_limits()
44    replace_stdout()
45    support.record_original_stdout(sys.stdout)
46
47    if ns.testdir:
48        # Prepend test directory to sys.path, so runtest() will be able
49        # to locate tests
50        sys.path.insert(0, os.path.abspath(ns.testdir))
51
52    # Some times __path__ and __file__ are not absolute (e.g. while running from
53    # Lib/) and, if we change the CWD to run the tests in a temporary dir, some
54    # imports might fail.  This affects only the modules imported before os.chdir().
55    # These modules are searched first in sys.path[0] (so '' -- the CWD) and if
56    # they are found in the CWD their __file__ and __path__ will be relative (this
57    # happens before the chdir).  All the modules imported after the chdir, are
58    # not found in the CWD, and since the other paths in sys.path[1:] are absolute
59    # (site.py absolutize them), the __file__ and __path__ will be absolute too.
60    # Therefore it is necessary to absolutize manually the __file__ and __path__ of
61    # the packages to prevent later imports to fail when the CWD is different.
62    for module in sys.modules.values():
63        if hasattr(module, '__path__'):
64            for index, path in enumerate(module.__path__):
65                module.__path__[index] = os.path.abspath(path)
66        if getattr(module, '__file__', None):
67            module.__file__ = os.path.abspath(module.__file__)
68
69    if ns.huntrleaks:
70        unittest.BaseTestSuite._cleanup = False
71
72    if ns.memlimit is not None:
73        support.set_memlimit(ns.memlimit)
74
75    if ns.threshold is not None:
76        gc.set_threshold(ns.threshold)
77
78    support.suppress_msvcrt_asserts(ns.verbose and ns.verbose >= 2)
79
80    support.use_resources = ns.use_resources
81
82    if hasattr(sys, 'addaudithook'):
83        # Add an auditing hook for all tests to ensure PySys_Audit is tested
84        def _test_audit_hook(name, args):
85            pass
86        sys.addaudithook(_test_audit_hook)
87
88    setup_unraisable_hook()
89    setup_threading_excepthook()
90
91    if ns.timeout is not None:
92        # For a slow buildbot worker, increase SHORT_TIMEOUT and LONG_TIMEOUT
93        support.SHORT_TIMEOUT = max(support.SHORT_TIMEOUT, ns.timeout / 40)
94        support.LONG_TIMEOUT = max(support.LONG_TIMEOUT, ns.timeout / 4)
95
96        # If --timeout is short: reduce timeouts
97        support.LOOPBACK_TIMEOUT = min(support.LOOPBACK_TIMEOUT, ns.timeout)
98        support.INTERNET_TIMEOUT = min(support.INTERNET_TIMEOUT, ns.timeout)
99        support.SHORT_TIMEOUT = min(support.SHORT_TIMEOUT, ns.timeout)
100        support.LONG_TIMEOUT = min(support.LONG_TIMEOUT, ns.timeout)
101
102    if ns.xmlpath:
103        from test.support.testresult import RegressionTestResult
104        RegressionTestResult.USE_XML = True
105
106    # Ensure there's a non-ASCII character in env vars at all times to force
107    # tests consider this case. See BPO-44647 for details.
108    if TESTFN_UNDECODABLE and os.supports_bytes_environ:
109        os.environb.setdefault(UNICODE_GUARD_ENV.encode(), TESTFN_UNDECODABLE)
110    elif FS_NONASCII:
111        os.environ.setdefault(UNICODE_GUARD_ENV, FS_NONASCII)
112
113
114def replace_stdout():
115    """Set stdout encoder error handler to backslashreplace (as stderr error
116    handler) to avoid UnicodeEncodeError when printing a traceback"""
117    stdout = sys.stdout
118    try:
119        fd = stdout.fileno()
120    except ValueError:
121        # On IDLE, sys.stdout has no file descriptor and is not a TextIOWrapper
122        # object. Leaving sys.stdout unchanged.
123        #
124        # Catch ValueError to catch io.UnsupportedOperation on TextIOBase
125        # and ValueError on a closed stream.
126        return
127
128    sys.stdout = open(fd, 'w',
129        encoding=stdout.encoding,
130        errors="backslashreplace",
131        closefd=False,
132        newline='\n')
133
134    def restore_stdout():
135        sys.stdout.close()
136        sys.stdout = stdout
137    atexit.register(restore_stdout)
138
139
140def _adjust_resource_limits():
141    """Adjust the system resource limits (ulimit) if needed."""
142    try:
143        import resource
144        from resource import RLIMIT_NOFILE, RLIM_INFINITY
145    except ImportError:
146        return
147    fd_limit, max_fds = resource.getrlimit(RLIMIT_NOFILE)
148    # On macOS the default fd limit is sometimes too low (256) for our
149    # test suite to succeed.  Raise it to something more reasonable.
150    # 1024 is a common Linux default.
151    desired_fds = 1024
152    if fd_limit < desired_fds and fd_limit < max_fds:
153        new_fd_limit = min(desired_fds, max_fds)
154        try:
155            resource.setrlimit(RLIMIT_NOFILE, (new_fd_limit, max_fds))
156            print(f"Raised RLIMIT_NOFILE: {fd_limit} -> {new_fd_limit}")
157        except (ValueError, OSError) as err:
158            print(f"Unable to raise RLIMIT_NOFILE from {fd_limit} to "
159                  f"{new_fd_limit}: {err}.")
160