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