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     if test_objects is None:
 165         from moztest.resolve import TestResolver
 166
 167         resolver = command_context._spawn(TestResolver)
 168         # If we were given test paths, try to find tests matching them.
 169         test_objects = resolver.resolve_tests(paths=tests, flavor="python")
 170     else:
 171         # We've received test_objects from |mach test|. We need to ignore
 172         # the subsuite because python-tests don't use this key like other
 173         # harnesses do and |mach test| doesn't realize this.
 174         subsuite = None
 175
 176     mp = TestManifest()
 177     mp.tests.extend(test_objects)
 178
 179     filters = []
 180     if subsuite == "default":
 181         filters.append(mpf.subsuite(None))
 182     elif subsuite:
 183         filters.append(mpf.subsuite(subsuite))
 184
 185     tests = mp.active_tests(filters=filters, disabled=False, python=3, **mozinfo.info)
 186
 187     if not tests:
 188         submsg = "for subsuite '{}' ".format(subsuite) if subsuite else ""
 189         message = (
 190             "TEST-UNEXPECTED-FAIL | No tests collected "
 191             + "{}(Not in PYTHON_UNITTEST_MANIFESTS?)".format(submsg)
 192         )
 193         command_context.log(logging.WARN, "python-test", {}, message)
 194         return 1
 195
 196     parallel = []
 197     sequential = []
 198     os.environ.setdefault("PYTEST_ADDOPTS", "")
 199
 200     if extra:
 201         os.environ["PYTEST_ADDOPTS"] += " " + " ".join(extra)
 202
 203     installed_requirements = set()
 204     for test in tests:
 205         if (
 206             test.get("requirements")
 207             and test["requirements"] not in installed_requirements
 208         ):
 209             command_context.virtualenv_manager.install_pip_requirements(
 210                 test["requirements"], quiet=True
 211             )
 212             installed_requirements.add(test["requirements"])
 213
 214     if exitfirst:
 215         sequential = tests
 216         os.environ["PYTEST_ADDOPTS"] += " -x"
 217     else:
 218         for test in tests:
 219             if test.get("sequential"):
 220                 sequential.append(test)
 221             else:
 222                 parallel.append(test)
 223
 224     jobs = jobs or cpu_count()
 225
 226     return_code = 0
 227     failure_output = []
 228
 229     def on_test_finished(result):
 230         output, ret, test_path = result
 231
 232         if ret:
 233             # Log the output of failed tests at the end so it's easy to find.
 234             failure_output.extend(output)
 235
 236             if not return_code:
 237                 command_context.log(
 238                     logging.ERROR,
 239                     "python-test",
 240                     {"test_path": test_path, "ret": ret},
 241                     "Setting retcode to {ret} from {test_path}",
 242                 )
 243         else:
 244             for line in output:
 245                 command_context.log(
 246                     logging.INFO, "python-test", {"line": line.rstrip()}, "{line}"
 247                 )
 248
 249         return return_code or ret
 250
 251     with tqdm(
 252         total=(len(parallel) + len(sequential)),
 253         unit="Test",
 254         desc="Tests Completed",
 255         initial=0,
 256     ) as progress_bar:
 257         try:
 258             with ThreadPoolExecutor(max_workers=jobs) as executor:
 259                 futures = []
 260
 261                 for test in parallel:
 262                     command_context.log(
 263                         logging.DEBUG,
 264                         "python-test",
 265                         {"line": f"Launching thread for test {test['file_relpath']}"},
 266                         "{line}",
 267                     )
 268                     futures.append(
 269                         executor.submit(
 270                             _run_python_test, command_context, test, jobs, verbose
 271                         )
 272                     )
 273
 274                 try:
 275                     for future in as_completed(futures):
 276                         progress_bar.clear()
 277                         return_code = on_test_finished(future.result())
 278                         progress_bar.update(1)
 279                 except KeyboardInterrupt:
 280                     # Hack to force stop currently running threads.
 281                     # https://gist.github.com/clchiou/f2608cbe54403edb0b13
 282                     executor._threads.clear()
 283                     thread._threads_queues.clear()
 284                     raise
 285
 286             for test in sequential:
 287                 test_result = _run_python_test(command_context, test, jobs, verbose)
 288
 289                 progress_bar.clear()
 290                 return_code = on_test_finished(test_result)
 291                 if return_code and exitfirst:
 292                     break
 293
 294                 progress_bar.update(1)
 295         finally:
 296             progress_bar.clear()
 297             # Now log all failures (even if there was a KeyboardInterrupt or other exception).
 298             for line in failure_output:
 299                 command_context.log(
 300                     logging.INFO, "python-test", {"line": line.rstrip()}, "{line}"
 301                 )
 302
 303     command_context.log(
 304         logging.INFO,
 305         "python-test",
 306         {"return_code": return_code},
 307         "Return code from mach python-test: {return_code}",
 308     )
 309
 310     return return_code
 311
 312
 313 def _run_python_test(command_context, test, jobs, verbose):
 314     output = []
 315
 316     def _log(line):
 317         # Buffer messages if more than one worker to avoid interleaving
 318         if jobs > 1:
 319             output.append(line)
 320         else:
 321             command_context.log(
 322                 logging.INFO, "python-test", {"line": line.rstrip()}, "{line}"
 323             )
 324
 325     _log(test["path"])
 326     python = command_context.virtualenv_manager.python_path
 327     cmd = [python, test["path"]]
 328     env = os.environ.copy()
 329     env["PYTHONDONTWRITEBYTECODE"] = "1"
 330
 331     result = subprocess.run(
 332         cmd,
 333         env=env,
 334         stdout=subprocess.PIPE,
 335         stderr=subprocess.STDOUT,
 336         universal_newlines=True,
 337         encoding="UTF-8",
 338     )
 339
 340     return_code = result.returncode
 341
 342     file_displayed_test = False
 343
 344     for line in result.stdout.split(os.linesep):
 345         if not file_displayed_test:
 346             test_ran = "Ran" in line or "collected" in line or line.startswith("TEST-")
 347             if test_ran:
 348                 file_displayed_test = True
 349
 350         # Hack to make sure treeherder highlights pytest failures
 351         if "FAILED" in line.rsplit(" ", 1)[-1]:
 352             line = line.replace("FAILED", "TEST-UNEXPECTED-FAIL")
 353
 354         _log(line)
 355
 356     if not file_displayed_test:
 357         return_code = 1
 358         _log(
 359             "TEST-UNEXPECTED-FAIL | No test output (missing mozunit.main() "
 360             "call?): {}".format(test["path"])
 361         )
 362
 363     if verbose:
 364         if return_code != 0:
 365             _log("Test failed: {}".format(test["path"]))
 366         else:
 367             _log("Test passed: {}".format(test["path"]))
 368
 369     return output, return_code, test["path"]