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