diff options
-rw-r--r-- | zephyr/zmake/tests/test_reexec.py | 59 | ||||
-rw-r--r-- | zephyr/zmake/zmake/__main__.py | 52 |
2 files changed, 111 insertions, 0 deletions
diff --git a/zephyr/zmake/tests/test_reexec.py b/zephyr/zmake/tests/test_reexec.py new file mode 100644 index 0000000000..9f25b5a834 --- /dev/null +++ b/zephyr/zmake/tests/test_reexec.py @@ -0,0 +1,59 @@ +# Copyright 2021 The Chromium OS Authors. All rights reserved. +# Use of this source code is governed by a BSD-style license that can be +# found in the LICENSE file. +"""Test the zmake re-exec functionality.""" + +import os +import sys +import unittest.mock as mock + +import pytest + +import zmake.__main__ as main + + +@pytest.fixture +def fake_env(monkeypatch): + environ = {} + monkeypatch.setattr(os, "environ", environ) + return environ + + +@pytest.fixture +def mock_execve(): + with mock.patch("os.execve", autospec=True) as mocked_function: + yield mocked_function + + +def test_out_of_chroot(fake_env, mock_execve): + # When CROS_WORKON_SRCROOT is not set, we should not re-exec. + main.maybe_reexec(["--help"]) + mock_execve.assert_not_called() + + +def test_pythonpath_set(fake_env, mock_execve): + # With PYTHONPATH set, we should not re-exec. + fake_env["CROS_WORKON_SRCROOT"] = "/mnt/host/source" + fake_env["PYTHONPATH"] = "/foo/bar/baz" + main.maybe_reexec(["--help"]) + mock_execve.assert_not_called() + + +def test_zmake_does_not_exist(fake_env, mock_execve): + # When zmake is not at src/platform/ec/zephyr/zmake, don't re-exec. + fake_env["CROS_WORKON_SRCROOT"] = "/this/does/not/exist" + main.maybe_reexec(["--help"]) + mock_execve.assert_not_called() + + +def test_zmake_reexec(fake_env, mock_execve): + # Nothing else applies? The re-exec should happen. + fake_env["CROS_WORKON_SRCROOT"] = "/mnt/host/source" + main.maybe_reexec(["--help"]) + new_env = dict(fake_env) + new_env["PYTHONPATH"] = "/mnt/host/source/src/platform/ec/zephyr/zmake" + mock_execve.assert_called_once_with( + sys.executable, + [sys.executable, "-m", "zmake", "--help"], + new_env, + ) diff --git a/zephyr/zmake/zmake/__main__.py b/zephyr/zmake/zmake/__main__.py index 74845d9bbc..a50ca3b08d 100644 --- a/zephyr/zmake/zmake/__main__.py +++ b/zephyr/zmake/zmake/__main__.py @@ -6,6 +6,7 @@ import argparse import inspect import logging +import os import pathlib import sys @@ -13,6 +14,55 @@ import zmake.multiproc as multiproc import zmake.zmake as zm +def maybe_reexec(argv): + """Re-exec zmake from the EC source tree, if possible and desired. + + Zmake installs into the users' chroot, which makes it convenient + to execute, but can sometimes become tedious when zmake changes + land and users haven't upgraded their chroots yet. + + We can partially subvert this problem by re-execing zmake from the + source if it's available. This won't make it so developers never + need to upgrade their chroots (e.g., a toolchain upgrade could + require chroot upgrades), but at least makes it slightly more + convenient for an average repo sync. + + Args: + argv: The argument list passed to the main function, not + including the executable path. + + Returns: + None, if the re-exec did not happen, or never returns if the + re-exec did happen. + """ + # We only re-exec if we are inside of a chroot (since if installed + # standalone using pip, there's already an "editable install" + # feature for that in pip.) + env = dict(os.environ) + srcroot = env.get("CROS_WORKON_SRCROOT") + if not srcroot: + return + + # If for some reason we decide to move zmake in the future, then + # we don't want to use the re-exec logic. + zmake_path = ( + pathlib.Path(srcroot) / "src" / "platform" / "ec" / "zephyr" / "zmake" + ).resolve() + if not zmake_path.is_dir(): + return + + # If PYTHONPATH is set, it is either because we just did a + # re-exec, or because the user wants to run a specific copy of + # zmake. In either case, we don't want to re-exec. + if "PYTHONPATH" in env: + return + + # Set PYTHONPATH so that we run zmake from source. + env["PYTHONPATH"] = str(zmake_path) + + os.execve(sys.executable, [sys.executable, "-m", "zmake", *argv], env) + + def call_with_namespace(func, namespace): """Call a function with arguments applied from a Namespace. @@ -55,6 +105,8 @@ def main(argv=None): if argv is None: argv = sys.argv[1:] + maybe_reexec(argv) + parser = argparse.ArgumentParser() parser.add_argument( "--checkout", type=pathlib.Path, help="Path to ChromiumOS checkout" |