python/mach_commands.py

   1 # This Source Code Form is subject to the terms of the Mozilla Public
   2 # License, v. 2.0. If a copy of the MPL was not distributed with this
   3 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
   4
   5 import argparse
   6 import logging
   7 import os
   8 import subprocess
   9 import tempfile
  10 from concurrent.futures import ThreadPoolExecutor, as_completed, thread
  11
  12 import mozinfo
  13 from mach.decorators import Command, CommandArgument
  14 from manifestparser import TestManifest
  15 from manifestparser import filters as mpf
  16 from mozbuild.util import cpu_count
  17 from mozfile import which
  18 from tqdm import tqdm
  19
  20
  21 @Command("python", category="devenv", description="Run Python.")
  22 @CommandArgument(
  23     "--exec-file", default=None, help="Execute this Python file using `exec`"
  24 )
  25 @CommandArgument(
  26     "--ipython",
  27     action="store_true",
  28     default=False,
  29     help="Use ipython instead of the default Python REPL.",
  30 )
  31 @CommandArgument(
  32     "--virtualenv",
  33     default=None,
  34     help="Prepare and use the virtualenv with the provided name. If not specified, "
  35     "then the Mach context is used instead.",
  36 )
  37 @CommandArgument("args", nargs=argparse.REMAINDER)
  38 def python(
  39     command_context,
  40     exec_file,
  41     ipython,
  42     virtualenv,
  43     args,
  44 ):
  45     # Avoid logging the command
  46     command_context.log_manager.terminal_handler.setLevel(logging.CRITICAL)
  47
  48     # Note: subprocess requires native strings in os.environ on Windows.
  49     append_env = {"PYTHONDONTWRITEBYTECODE": str("1")}
  50
  51     if virtualenv:
  52         command_context._virtualenv_name = virtualenv
  53
  54     if exec_file:
  55         command_context.activate_virtualenv()
  56         exec(open(exec_file).read())
  57         return 0
  58
  59     if ipython:
  60         if virtualenv:
  61             command_context.virtualenv_manager.ensure()
  62             python_path = which(
  63                 "ipython", path=command_context.virtualenv_manager.bin_path
  64             )
  65             if not python_path:
  66                 raise Exception(
  67                     "--ipython was specified, but the provided "
  68                     '--virtualenv doesn\'t have "ipython" installed.'
  69                 )
  70         else:
  71             command_context._virtualenv_name = "ipython"
  72             command_context.virtualenv_manager.ensure()
  73             python_path = which(
  74                 "ipython", path=command_context.virtualenv_manager.bin_path
  75             )
  76     else:
  77         command_context.virtualenv_manager.ensure()
  78         python_path = command_context.virtualenv_manager.python_path
  79
  80     return command_context.run_process(
  81         [python_path] + args,
  82         pass_thru=True,  # Allow user to run Python interactively.
  83         ensure_exit_code=False,  # Don't throw on non-zero exit code.
  84         python_unbuffered=False,  # Leave input buffered.
  85         append_env=append_env,
  86     )
  87
  88
  89 @Command(
  90     "python-test",
  91     category="testing",
  92     virtualenv_name="python-test",
  93     description="Run Python unit tests with pytest.",
  94 )
  95 @CommandArgument(
  96     "-v", "--verbose", default=False, action="store_true", help="Verbose output."
  97 )
  98 @CommandArgument(
  99     "-j",
 100     "--jobs",
 101     default=None,
 102     type=int,
 103     help="Number of concurrent jobs to run. Default is the number of CPUs "
 104     "in the system.",
 105 )
 106 @CommandArgument(
 107     "-x",
 108     "--exitfirst",
 109     default=False,
 110     action="store_true",
 111     help="Runs all tests sequentially and breaks at the first failure.",
 112 )
 113 @CommandArgument(
 114     "--subsuite",
 115     default=None,
 116     help=(
 117         "Python subsuite to run. If not specified, all subsuites are run. "
 118         "Use the string `default` to only run tests without a subsuite."
 119     ),
 120 )
 121 @CommandArgument(
 122     "tests",
 123     nargs="*",
 124     metavar="TEST",
 125     help=(
 126         "Tests to run. Each test can be a single file or a directory. "
 127         "Default test resolution relies on PYTHON_UNITTEST_MANIFESTS."
 128     ),
 129 )
 130 @CommandArgument(
 131     "extra",
 132     nargs=argparse.REMAINDER,
 133     metavar="PYTEST ARGS",
 134     help=(
 135         "Arguments that aren't recognized by mach. These will be "
 136         "passed as it is to pytest"
 137     ),
 138 )
 139 def python_test(command_context, *args, **kwargs):
 140     try:
 141         tempdir = str(tempfile.mkdtemp(suffix="-python-test"))
 142         os.environ["PYTHON_TEST_TMP"] = tempdir
 143         return run_python_tests(command_context, *args, **kwargs)
 144     finally:
 145         import mozfile
 146
 147         mozfile.remove(tempdir)
 148
 149
 150 def run_python_tests(
 151     command_context,
 152     tests=None,
 153     test_objects=None,
 154     subsuite=None,
 155     verbose=False,
 156     jobs=None,
 157     exitfirst=False,
 158     extra=None,
 159     **kwargs,
 160 ):
 161     if test_objects is None:
 162         from moztest.resolve import TestResolver
 163
 164         resolver = command_context._spawn(TestResolver)
 165         # If we were given test paths, try to find tests matching them.
 166         test_objects = resolver.resolve_tests(paths=tests, flavor="python")
 167     else:
 168         # We've received test_objects from |mach test|. We need to ignore
 169         # the subsuite because python-tests don't use this key like other
 170         # harnesses do and |mach test| doesn't realize this.
 171         subsuite = None
 172
 173     mp = TestManifest()
 174     mp.tests.extend(test_objects)
 175
 176     filters = []
 177     if subsuite == "default":
 178         filters.append(mpf.subsuite(None))
 179     elif subsuite:
 180         filters.append(mpf.subsuite(subsuite))
 181
 182     tests = mp.active_tests(filters=filters, disabled=False, python=3, **mozinfo.info)
 183
 184     if not tests:
 185         submsg = "for subsuite '{}' ".format(subsuite) if subsuite else ""
 186         message = (
 187             "TEST-UNEXPECTED-FAIL | No tests collected "
 188             + "{}(Not in PYTHON_UNITTEST_MANIFESTS?)".format(submsg)
 189         )
 190         command_context.log(logging.WARN, "python-test", {}, message)
 191         return 1
 192
 193     parallel = []
 194     sequential = []
 195     os.environ.setdefault("PYTEST_ADDOPTS", "")
 196
 197     if extra:
 198         os.environ["PYTEST_ADDOPTS"] += " " + " ".join(extra)
 199
 200     installed_requirements = set()
 201     for test in tests:
 202         if (
 203             test.get("requirements")
 204             and test["requirements"] not in installed_requirements
 205         ):
 206             command_context.virtualenv_manager.install_pip_requirements(
 207                 test["requirements"], quiet=True
 208             )
 209             installed_requirements.add(test["requirements"])
 210
 211     if exitfirst:
 212         sequential = tests
 213         os.environ["PYTEST_ADDOPTS"] += " -x"
 214     else:
 215         for test in tests:
 216             if test.get("sequential"):
 217                 sequential.append(test)
 218             else:
 219                 parallel.append(test)
 220
 221     jobs = jobs or cpu_count()
 222
 223     return_code = 0
 224     failure_output = []
 225
 226     def on_test_finished(result):
 227         output, ret, test_path = result
 228
 229         if ret:
 230             # Log the output of failed tests at the end so it's easy to find.
 231             failure_output.extend(output)
 232
 233             if not return_code:
 234                 command_context.log(
 235                     logging.ERROR,
 236                     "python-test",
 237                     {"test_path": test_path, "ret": ret},
 238                     "Setting retcode to {ret} from {test_path}",
 239                 )
 240         else:
 241             for line in output:
 242                 command_context.log(
 243                     logging.INFO, "python-test", {"line": line.rstrip()}, "{line}"
 244                 )
 245
 246         return return_code or ret
 247
 248     with tqdm(
 249         total=(len(parallel) + len(sequential)),
 250         unit="Test",
 251         desc="Tests Completed",
 252         initial=0,
 253     ) as progress_bar:
 254         try:
 255             with ThreadPoolExecutor(max_workers=jobs) as executor:
 256                 futures = []
 257
 258                 for test in parallel:
 259                     command_context.log(
 260                         logging.DEBUG,
 261                         "python-test",
 262                         {"line": f"Launching thread for test {test['file_relpath']}"},
 263                         "{line}",
 264                     )
 265                     futures.append(
 266                         executor.submit(
 267                             _run_python_test, command_context, test, jobs, verbose
 268                         )
 269                     )
 270
 271                 try:
 272                     for future in as_completed(futures):
 273                         progress_bar.clear()
 274                         return_code = on_test_finished(future.result())
 275                         progress_bar.update(1)
 276                 except KeyboardInterrupt:
 277                     # Hack to force stop currently running threads.
 278                     # https://gist.github.com/clchiou/f2608cbe54403edb0b13
 279                     executor._threads.clear()
 280                     thread._threads_queues.clear()
 281                     raise
 282
 283             for test in sequential:
 284                 test_result = _run_python_test(command_context, test, jobs, verbose)
 285
 286                 progress_bar.clear()
 287                 return_code = on_test_finished(test_result)
 288                 if return_code and exitfirst:
 289                     break
 290
 291                 progress_bar.update(1)
 292         finally:
 293             progress_bar.clear()
 294             # Now log all failures (even if there was a KeyboardInterrupt or other exception).
 295             for line in failure_output:
 296                 command_context.log(
 297                     logging.INFO, "python-test", {"line": line.rstrip()}, "{line}"
 298                 )
 299
 300     command_context.log(
 301         logging.INFO,
 302         "python-test",
 303         {"return_code": return_code},
 304         "Return code from mach python-test: {return_code}",
 305     )
 306
 307     return return_code
 308
 309
 310 def _run_python_test(command_context, test, jobs, verbose):
 311     output = []
 312
 313     def _log(line):
 314         # Buffer messages if more than one worker to avoid interleaving
 315         if jobs > 1:
 316             output.append(line)
 317         else:
 318             command_context.log(
 319                 logging.INFO, "python-test", {"line": line.rstrip()}, "{line}"
 320             )
 321
 322     _log(test["path"])
 323     python = command_context.virtualenv_manager.python_path
 324     cmd = [python, test["path"]]
 325     env = os.environ.copy()
 326     env["PYTHONDONTWRITEBYTECODE"] = "1"
 327
 328     result = subprocess.run(
 329         cmd,
 330         env=env,
 331         stdout=subprocess.PIPE,
 332         stderr=subprocess.STDOUT,
 333         universal_newlines=True,
 334         encoding="UTF-8",
 335     )
 336
 337     return_code = result.returncode
 338
 339     file_displayed_test = False
 340
 341     for line in result.stdout.split(os.linesep):
 342         if not file_displayed_test:
 343             test_ran = "Ran" in line or "collected" in line or line.startswith("TEST-")
 344             if test_ran:
 345                 file_displayed_test = True
 346
 347         # Hack to make sure treeherder highlights pytest failures
 348         if "FAILED" in line.rsplit(" ", 1)[-1]:
 349             line = line.replace("FAILED", "TEST-UNEXPECTED-FAIL")
 350
 351         _log(line)
 352
 353     if not file_displayed_test:
 354         return_code = 1
 355         _log(
 356             "TEST-UNEXPECTED-FAIL | No test output (missing mozunit.main() "
 357             "call?): {}".format(test["path"])
 358         )
 359
 360     if verbose:
 361         if return_code != 0:
 362             _log("Test failed: {}".format(test["path"]))
 363         else:
 364             _log("Test passed: {}".format(test["path"]))
 365
 366     return output, return_code, test["path"]