afdo_tools: add chromeos-tree flag to update_kernel_afdo.py

...While this tool doesn't need the `repo` command, `gerrit` needs to be
invoked from within a ChromeOS tree. Allow that to be specified (or
autodetected, as necessary).

BUG=b:329449239
TEST=Ran locally.

Change-Id: Icfc2fd19d99c07ea30d7ef8c7ac11770a790a867
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/5377048
Reviewed-by: Jordan Abrahams-Whitehead <ajordanr@google.com>
Tested-by: George Burgess <gbiv@chromium.org>
Commit-Queue: George Burgess <gbiv@chromium.org>

2 files changed, 48 insertions, 1 deletions

diff --git a/afdo_tools/update_kernel_afdo.py b/afdo_tools/update_kernel_afdo.py
index 76e336e4..d91bd391 100755
--- a/afdo_tools/update_kernel_afdo.py
+++ b/afdo_tools/update_kernel_afdo.py
@@ -139,6 +139,15 @@ def get_parser():
         """,
     )
     parser.add_argument(
+        "--chromeos-tree",
+        type=Path,
+        help="""
+        Root of a ChromeOS tree. This is optional to pass in, but doing so
+        unlocks extra convenience features on `--upload`. This script will try
+        to autodetect a tree if this isn't specified.
+        """,
+    )
+    parser.add_argument(
         "channel",
         nargs="*",
         type=Channel.parse,
@@ -729,6 +738,7 @@ def parse_cl_from_upload_output(upload_output: str) -> str:
 
 def upload_head_to_gerrit(
     toolchain_utils: Path,
+    chromeos_tree: Optional[Path],
     branch: GitBranch,
 ):
     """Uploads HEAD to gerrit as a CL, and sets reviewers/CCs."""
@@ -762,6 +772,13 @@ def upload_head_to_gerrit(
     cl_id = parse_cl_from_upload_output(run_result.stdout)
     logging.info("Uploaded CL http://crrev.com/c/%s successfully.", cl_id)
 
+    if chromeos_tree is None:
+        logging.info(
+            "Skipping gerrit convenience commands, since no CrOS tree was "
+            "specified."
+        )
+        return
+
     # To make the life of the reviewers marginally easier, click buttons
     # automatically.
     gerrit_commands = (
@@ -776,6 +793,7 @@ def upload_head_to_gerrit(
         # script is expeted to be used.
         return_code = subprocess.run(
             cmd,
+            cwd=chromeos_tree,
             check=False,
             stdin=subprocess.DEVNULL,
         ).returncode
@@ -786,6 +804,13 @@ def upload_head_to_gerrit(
             )
 
 
+def find_chromeos_tree_root(a_dir: Path) -> Optional[Path]:
+    for parent in a_dir.parents:
+        if (parent / ".repo").is_dir():
+            return parent
+    return None
+
+
 def main(argv: List[str]) -> None:
     my_dir = Path(__file__).resolve().parent
     toolchain_utils = my_dir.parent
@@ -797,6 +822,12 @@ def main(argv: List[str]) -> None:
         level=logging.DEBUG if opts.debug else logging.INFO,
     )
 
+    chromeos_tree = opts.chromeos_tree
+    if not chromeos_tree:
+        chromeos_tree = find_chromeos_tree_root(my_dir)
+        if chromeos_tree:
+            logging.info("Autodetected ChromeOS tree root at %s", chromeos_tree)
+
     if opts.fetch:
         logging.info("Fetching in %s...", toolchain_utils)
         git_fetch(toolchain_utils)
@@ -832,7 +863,7 @@ def main(argv: List[str]) -> None:
             commit_new_profiles(worktree, channel, result.had_failures)
             if opts.upload:
                 logging.info("New profiles were committed. Uploading...")
-                upload_head_to_gerrit(worktree, branch)
+                upload_head_to_gerrit(worktree, chromeos_tree, branch)
             else:
                 logging.info(
                     "--upload not specified. Leaving commit for %s at %s",
diff --git a/afdo_tools/update_kernel_afdo_test.py b/afdo_tools/update_kernel_afdo_test.py
index 60f52138..79ed9d1c 100755
--- a/afdo_tools/update_kernel_afdo_test.py
+++ b/afdo_tools/update_kernel_afdo_test.py
@@ -337,6 +337,22 @@ TOTAL: 2 objects, 1234 bytes (1.1KiB)
                 GERRIT_OUTPUT_WITH_TWO_CLS
             )
 
+    def test_repo_autodetects_nothing_if_no_repo_dir(self):
+        self.assertIsNone(
+            update_kernel_afdo.find_chromeos_tree_root(
+                Path("/does/not/exist/nor/is/under/a/repo")
+            )
+        )
+
+    def test_repo_autodetects_repo_dir_correctly(self):
+        tmpdir = self.make_tempdir()
+        test_subdir = tmpdir / "a/directory/and/another/one"
+        test_subdir.mkdir(parents=True)
+        (tmpdir / ".repo").mkdir()
+        self.assertEqual(
+            tmpdir, update_kernel_afdo.find_chromeos_tree_root(test_subdir)
+        )
+
 
 if __name__ == "__main__":
     unittest.main()