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 from multiprocessing import cpu_count
  12
  13 import mozinfo
  14 from mach.decorators import Command, CommandArgument
  15 from manifestparser import TestManifest
  16 from manifestparser import filters as mpf
  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"]