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