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