michael@0: # This Source Code Form is subject to the terms of the Mozilla Public
michael@0: # License, v. 2.0. If a copy of the MPL was not distributed with this
michael@0: # file, You can obtain one at http://mozilla.org/MPL/2.0/.
michael@0: 
michael@0: from __future__ import print_function, unicode_literals
michael@0: 
michael@0: import argparse
michael@0: import glob
michael@0: import logging
michael@0: import mozpack.path
michael@0: import os
michael@0: import sys
michael@0: 
michael@0: from mozbuild.base import (
michael@0:     MachCommandBase,
michael@0: )
michael@0: 
michael@0: from mach.decorators import (
michael@0:     CommandArgument,
michael@0:     CommandProvider,
michael@0:     Command,
michael@0: )
michael@0: 
michael@0: 
michael@0: @CommandProvider
michael@0: class MachCommands(MachCommandBase):
michael@0:     @Command('python', category='devenv',
michael@0:         allow_all_args=True,
michael@0:         description='Run Python.')
michael@0:     @CommandArgument('args', nargs=argparse.REMAINDER)
michael@0:     def python(self, args):
michael@0:         # Avoid logging the command
michael@0:         self.log_manager.terminal_handler.setLevel(logging.CRITICAL)
michael@0: 
michael@0:         self._activate_virtualenv()
michael@0: 
michael@0:         return self.run_process([self.virtualenv_manager.python_path] + args,
michael@0:             pass_thru=True, # Allow user to run Python interactively.
michael@0:             ensure_exit_code=False, # Don't throw on non-zero exit code.
michael@0:             # Note: subprocess requires native strings in os.environ on Windows
michael@0:             append_env={b'PYTHONDONTWRITEBYTECODE': str('1')})
michael@0: 
michael@0:     @Command('python-test', category='testing',
michael@0:         description='Run Python unit tests.')
michael@0:     @CommandArgument('--verbose',
michael@0:         default=False,
michael@0:         action='store_true',
michael@0:         help='Verbose output.')
michael@0:     @CommandArgument('--stop',
michael@0:         default=False,
michael@0:         action='store_true',
michael@0:         help='Stop running tests after the first error or failure.')
michael@0:     @CommandArgument('tests', nargs='+',
michael@0:         metavar='TEST',
michael@0:         help='Tests to run. Each test can be a single file or a directory.')
michael@0:     def python_test(self, tests, verbose=False, stop=False):
michael@0:         self._activate_virtualenv()
michael@0: 
michael@0:         # Python's unittest, and in particular discover, has problems with
michael@0:         # clashing namespaces when importing multiple test modules. What follows
michael@0:         # is a simple way to keep environments separate, at the price of
michael@0:         # launching Python multiple times. This also runs tests via mozunit,
michael@0:         # which produces output in the format Mozilla infrastructure expects.
michael@0:         return_code = 0
michael@0:         files = []
michael@0:         for test in tests:
michael@0:             if test.endswith('.py') and os.path.isfile(test):
michael@0:                 files.append(test)
michael@0:             elif os.path.isfile(test + '.py'):
michael@0:                 files.append(test + '.py')
michael@0:             elif os.path.isdir(test):
michael@0:                 files += glob.glob(mozpack.path.join(test, 'test*.py'))
michael@0:                 files += glob.glob(mozpack.path.join(test, 'unit*.py'))
michael@0:             else:
michael@0:                 self.log(logging.WARN, 'python-test', {'test': test},
michael@0:                          'TEST-UNEXPECTED-FAIL | Invalid test: {test}')
michael@0:                 if stop:
michael@0:                     return 1
michael@0: 
michael@0:         for f in files:
michael@0:             file_displayed_test = [] # Used as a boolean.
michael@0:             def _line_handler(line):
michael@0:                 if not file_displayed_test and line.startswith('TEST-'):
michael@0:                     file_displayed_test.append(True)
michael@0: 
michael@0:             inner_return_code = self.run_process(
michael@0:                 [self.virtualenv_manager.python_path, f],
michael@0:                 ensure_exit_code=False, # Don't throw on non-zero exit code.
michael@0:                 log_name='python-test',
michael@0:                 # subprocess requires native strings in os.environ on Windows
michael@0:                 append_env={b'PYTHONDONTWRITEBYTECODE': str('1')},
michael@0:                 line_handler=_line_handler)
michael@0:             return_code += inner_return_code
michael@0: 
michael@0:             if not file_displayed_test:
michael@0:                 self.log(logging.WARN, 'python-test', {'file': f},
michael@0:                          'TEST-UNEXPECTED-FAIL | No test output (missing mozunit.main() call?): {file}')
michael@0: 
michael@0:             if verbose:
michael@0:                 if inner_return_code != 0:
michael@0:                     self.log(logging.INFO, 'python-test', {'file': f},
michael@0:                              'Test failed: {file}')
michael@0:                 else:
michael@0:                     self.log(logging.INFO, 'python-test', {'file': f},
michael@0:                              'Test passed: {file}')
michael@0:             if stop and return_code > 0:
michael@0:                 return 1
michael@0: 
michael@0:         return 0 if return_code == 0 else 1