summaryrefslogtreecommitdiffstats
path: root/python/mach_commands.py
diff options
context:
space:
mode:
Diffstat (limited to 'python/mach_commands.py')
-rw-r--r--python/mach_commands.py366
1 files changed, 366 insertions, 0 deletions
diff --git a/python/mach_commands.py b/python/mach_commands.py
new file mode 100644
index 0000000000..d4f1f67efe
--- /dev/null
+++ b/python/mach_commands.py
@@ -0,0 +1,366 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+import argparse
+import logging
+import os
+import subprocess
+import tempfile
+from concurrent.futures import ThreadPoolExecutor, as_completed, thread
+from multiprocessing import cpu_count
+
+import mozinfo
+from mach.decorators import Command, CommandArgument
+from manifestparser import TestManifest
+from manifestparser import filters as mpf
+from mozfile import which
+from tqdm import tqdm
+
+
+@Command("python", category="devenv", description="Run Python.")
+@CommandArgument(
+ "--exec-file", default=None, help="Execute this Python file using `exec`"
+)
+@CommandArgument(
+ "--ipython",
+ action="store_true",
+ default=False,
+ help="Use ipython instead of the default Python REPL.",
+)
+@CommandArgument(
+ "--virtualenv",
+ default=None,
+ help="Prepare and use the virtualenv with the provided name. If not specified, "
+ "then the Mach context is used instead.",
+)
+@CommandArgument("args", nargs=argparse.REMAINDER)
+def python(
+ command_context,
+ exec_file,
+ ipython,
+ virtualenv,
+ args,
+):
+ # Avoid logging the command
+ command_context.log_manager.terminal_handler.setLevel(logging.CRITICAL)
+
+ # Note: subprocess requires native strings in os.environ on Windows.
+ append_env = {"PYTHONDONTWRITEBYTECODE": str("1")}
+
+ if virtualenv:
+ command_context._virtualenv_name = virtualenv
+
+ if exec_file:
+ command_context.activate_virtualenv()
+ exec(open(exec_file).read())
+ return 0
+
+ if ipython:
+ if virtualenv:
+ command_context.virtualenv_manager.ensure()
+ python_path = which(
+ "ipython", path=command_context.virtualenv_manager.bin_path
+ )
+ if not python_path:
+ raise Exception(
+ "--ipython was specified, but the provided "
+ '--virtualenv doesn\'t have "ipython" installed.'
+ )
+ else:
+ command_context._virtualenv_name = "ipython"
+ command_context.virtualenv_manager.ensure()
+ python_path = which(
+ "ipython", path=command_context.virtualenv_manager.bin_path
+ )
+ else:
+ command_context.virtualenv_manager.ensure()
+ python_path = command_context.virtualenv_manager.python_path
+
+ return command_context.run_process(
+ [python_path] + args,
+ pass_thru=True, # Allow user to run Python interactively.
+ ensure_exit_code=False, # Don't throw on non-zero exit code.
+ python_unbuffered=False, # Leave input buffered.
+ append_env=append_env,
+ )
+
+
+@Command(
+ "python-test",
+ category="testing",
+ virtualenv_name="python-test",
+ description="Run Python unit tests with pytest.",
+)
+@CommandArgument(
+ "-v", "--verbose", default=False, action="store_true", help="Verbose output."
+)
+@CommandArgument(
+ "-j",
+ "--jobs",
+ default=None,
+ type=int,
+ help="Number of concurrent jobs to run. Default is the number of CPUs "
+ "in the system.",
+)
+@CommandArgument(
+ "-x",
+ "--exitfirst",
+ default=False,
+ action="store_true",
+ help="Runs all tests sequentially and breaks at the first failure.",
+)
+@CommandArgument(
+ "--subsuite",
+ default=None,
+ help=(
+ "Python subsuite to run. If not specified, all subsuites are run. "
+ "Use the string `default` to only run tests without a subsuite."
+ ),
+)
+@CommandArgument(
+ "tests",
+ nargs="*",
+ metavar="TEST",
+ help=(
+ "Tests to run. Each test can be a single file or a directory. "
+ "Default test resolution relies on PYTHON_UNITTEST_MANIFESTS."
+ ),
+)
+@CommandArgument(
+ "extra",
+ nargs=argparse.REMAINDER,
+ metavar="PYTEST ARGS",
+ help=(
+ "Arguments that aren't recognized by mach. These will be "
+ "passed as it is to pytest"
+ ),
+)
+def python_test(command_context, *args, **kwargs):
+ try:
+ tempdir = str(tempfile.mkdtemp(suffix="-python-test"))
+ os.environ["PYTHON_TEST_TMP"] = tempdir
+ return run_python_tests(command_context, *args, **kwargs)
+ finally:
+ import mozfile
+
+ mozfile.remove(tempdir)
+
+
+def run_python_tests(
+ command_context,
+ tests=None,
+ test_objects=None,
+ subsuite=None,
+ verbose=False,
+ jobs=None,
+ exitfirst=False,
+ extra=None,
+ **kwargs,
+):
+ if test_objects is None:
+ from moztest.resolve import TestResolver
+
+ resolver = command_context._spawn(TestResolver)
+ # If we were given test paths, try to find tests matching them.
+ test_objects = resolver.resolve_tests(paths=tests, flavor="python")
+ else:
+ # We've received test_objects from |mach test|. We need to ignore
+ # the subsuite because python-tests don't use this key like other
+ # harnesses do and |mach test| doesn't realize this.
+ subsuite = None
+
+ mp = TestManifest()
+ mp.tests.extend(test_objects)
+
+ filters = []
+ if subsuite == "default":
+ filters.append(mpf.subsuite(None))
+ elif subsuite:
+ filters.append(mpf.subsuite(subsuite))
+
+ tests = mp.active_tests(filters=filters, disabled=False, python=3, **mozinfo.info)
+
+ if not tests:
+ submsg = "for subsuite '{}' ".format(subsuite) if subsuite else ""
+ message = (
+ "TEST-UNEXPECTED-FAIL | No tests collected "
+ + "{}(Not in PYTHON_UNITTEST_MANIFESTS?)".format(submsg)
+ )
+ command_context.log(logging.WARN, "python-test", {}, message)
+ return 1
+
+ parallel = []
+ sequential = []
+ os.environ.setdefault("PYTEST_ADDOPTS", "")
+
+ if extra:
+ os.environ["PYTEST_ADDOPTS"] += " " + " ".join(extra)
+
+ installed_requirements = set()
+ for test in tests:
+ if (
+ test.get("requirements")
+ and test["requirements"] not in installed_requirements
+ ):
+ command_context.virtualenv_manager.install_pip_requirements(
+ test["requirements"], quiet=True
+ )
+ installed_requirements.add(test["requirements"])
+
+ if exitfirst:
+ sequential = tests
+ os.environ["PYTEST_ADDOPTS"] += " -x"
+ else:
+ for test in tests:
+ if test.get("sequential"):
+ sequential.append(test)
+ else:
+ parallel.append(test)
+
+ jobs = jobs or cpu_count()
+
+ return_code = 0
+ failure_output = []
+
+ def on_test_finished(result):
+ output, ret, test_path = result
+
+ if ret:
+ # Log the output of failed tests at the end so it's easy to find.
+ failure_output.extend(output)
+
+ if not return_code:
+ command_context.log(
+ logging.ERROR,
+ "python-test",
+ {"test_path": test_path, "ret": ret},
+ "Setting retcode to {ret} from {test_path}",
+ )
+ else:
+ for line in output:
+ command_context.log(
+ logging.INFO, "python-test", {"line": line.rstrip()}, "{line}"
+ )
+
+ return return_code or ret
+
+ with tqdm(
+ total=(len(parallel) + len(sequential)),
+ unit="Test",
+ desc="Tests Completed",
+ initial=0,
+ ) as progress_bar:
+ try:
+ with ThreadPoolExecutor(max_workers=jobs) as executor:
+ futures = []
+
+ for test in parallel:
+ command_context.log(
+ logging.DEBUG,
+ "python-test",
+ {"line": f"Launching thread for test {test['file_relpath']}"},
+ "{line}",
+ )
+ futures.append(
+ executor.submit(
+ _run_python_test, command_context, test, jobs, verbose
+ )
+ )
+
+ try:
+ for future in as_completed(futures):
+ progress_bar.clear()
+ return_code = on_test_finished(future.result())
+ progress_bar.update(1)
+ except KeyboardInterrupt:
+ # Hack to force stop currently running threads.
+ # https://gist.github.com/clchiou/f2608cbe54403edb0b13
+ executor._threads.clear()
+ thread._threads_queues.clear()
+ raise
+
+ for test in sequential:
+ test_result = _run_python_test(command_context, test, jobs, verbose)
+
+ progress_bar.clear()
+ return_code = on_test_finished(test_result)
+ if return_code and exitfirst:
+ break
+
+ progress_bar.update(1)
+ finally:
+ progress_bar.clear()
+ # Now log all failures (even if there was a KeyboardInterrupt or other exception).
+ for line in failure_output:
+ command_context.log(
+ logging.INFO, "python-test", {"line": line.rstrip()}, "{line}"
+ )
+
+ command_context.log(
+ logging.INFO,
+ "python-test",
+ {"return_code": return_code},
+ "Return code from mach python-test: {return_code}",
+ )
+
+ return return_code
+
+
+def _run_python_test(command_context, test, jobs, verbose):
+ output = []
+
+ def _log(line):
+ # Buffer messages if more than one worker to avoid interleaving
+ if jobs > 1:
+ output.append(line)
+ else:
+ command_context.log(
+ logging.INFO, "python-test", {"line": line.rstrip()}, "{line}"
+ )
+
+ _log(test["path"])
+ python = command_context.virtualenv_manager.python_path
+ cmd = [python, test["path"]]
+ env = os.environ.copy()
+ env["PYTHONDONTWRITEBYTECODE"] = "1"
+
+ result = subprocess.run(
+ cmd,
+ env=env,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ universal_newlines=True,
+ encoding="UTF-8",
+ )
+
+ return_code = result.returncode
+
+ file_displayed_test = False
+
+ for line in result.stdout.split(os.linesep):
+ if not file_displayed_test:
+ test_ran = "Ran" in line or "collected" in line or line.startswith("TEST-")
+ if test_ran:
+ file_displayed_test = True
+
+ # Hack to make sure treeherder highlights pytest failures
+ if "FAILED" in line.rsplit(" ", 1)[-1]:
+ line = line.replace("FAILED", "TEST-UNEXPECTED-FAIL")
+
+ _log(line)
+
+ if not file_displayed_test:
+ return_code = 1
+ _log(
+ "TEST-UNEXPECTED-FAIL | No test output (missing mozunit.main() "
+ "call?): {}".format(test["path"])
+ )
+
+ if verbose:
+ if return_code != 0:
+ _log("Test failed: {}".format(test["path"]))
+ else:
+ _log("Test passed: {}".format(test["path"]))
+
+ return output, return_code, test["path"]