llvm_tools: fix nightly_revert_checker cherrypick logic

This param should be called `chromeos_tree`. While refactoring that, I
saw that `find_chroot` wasn't a great name for "get the chromeos tree
(or android tree (unused)), interesting SHAs, and email recipients."
Moving that logic into main seems simplest.

BUG=b:333720269
TEST=repo upload

Change-Id: I0b4b0f06a839553f4bad9c8fbcd0b9a209e71f24
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/5507516
Tested-by: George Burgess <gbiv@chromium.org>
Reviewed-by: Ryan Beltran <ryanbeltran@chromium.org>
Commit-Queue: George Burgess <gbiv@chromium.org>

2 files changed, 24 insertions, 35 deletions

diff --git a/llvm_tools/nightly_revert_checker.py b/llvm_tools/nightly_revert_checker.py
index 2a79cfd0..284ba8ae 100644
--- a/llvm_tools/nightly_revert_checker.py
+++ b/llvm_tools/nightly_revert_checker.py
@@ -16,7 +16,6 @@ import os
 from pathlib import Path
 import pprint
 import subprocess
-import sys
 import time
 from typing import Any, Callable, Dict, List, NamedTuple, Set, Tuple
 
@@ -311,7 +310,7 @@ def locate_new_reverts_across_shas(
 
 
 def do_cherrypick(
-    chroot_path: str,
+    chromeos_path: str,
     llvm_dir: str,
     repository: str,
     interesting_shas: List[Tuple[str, str]],
@@ -338,7 +337,7 @@ def do_cherrypick(
             # We upload reverts for all platforms by default, since there's
             # no real reason for them to be CrOS-specific.
             get_upstream_patch.get_from_upstream(
-                chroot_path=chroot_path,
+                chromeos_path=chromeos_path,
                 create_cl=True,
                 start_sha=reverted_sha,
                 patches=[sha],
@@ -553,35 +552,6 @@ def parse_args(argv: List[str]) -> argparse.Namespace:
     return parser.parse_args(argv)
 
 
-def find_chroot(
-    opts: argparse.Namespace, cc: List[str]
-) -> Tuple[str, List[Tuple[str, str]], _EmailRecipients]:
-    if opts.repository == "chromeos":
-        chroot_path = opts.chromeos_dir
-        return (
-            chroot_path,
-            _find_interesting_chromeos_shas(chroot_path),
-            _EmailRecipients(well_known=["mage"], direct=cc),
-        )
-    elif opts.repository == "android":
-        if opts.action == "cherry-pick":
-            raise RuntimeError(
-                "android doesn't currently support automatic cherry-picking."
-            )
-
-        chroot_path = opts.android_llvm_toolchain_dir
-        return (
-            chroot_path,
-            _find_interesting_android_shas(chroot_path),
-            _EmailRecipients(
-                well_known=[],
-                direct=["android-llvm-dev@google.com"] + cc,
-            ),
-        )
-    else:
-        raise ValueError(f"Unknown repository {opts.repository}")
-
-
 def main(argv: List[str]) -> int:
     opts = parse_args(argv)
 
@@ -598,7 +568,21 @@ def main(argv: List[str]) -> int:
     reviewers = opts.reviewers if opts.reviewers else []
     cc = opts.cc if opts.cc else []
 
-    chroot_path, interesting_shas, recipients = find_chroot(opts, cc)
+    if opts.repository == "chromeos":
+        chromeos_path = opts.chromeos_path
+        interesting_shas = _find_interesting_chromeos_shas(chromeos_path)
+        recipients = _EmailRecipients(well_known=["mage"], direct=cc)
+    elif opts.repository == "android":
+        interesting_shas = _find_interesting_android_shas(
+            opts.android_llvm_toolchain_dir
+        )
+        recipients = _EmailRecipients(
+            well_known=[],
+            direct=["android-llvm-dev@google.com"] + cc,
+        )
+    else:
+        raise ValueError(f"Unknown repository {opts.repository}")
+
     logging.info("Interesting SHAs were %r", interesting_shas)
 
     state = _read_state(state_file)
@@ -607,8 +591,13 @@ def main(argv: List[str]) -> int:
     # We want to be as free of obvious side-effects as possible in case
     # something above breaks. Hence, action as late as possible.
     if action == "cherry-pick":
+        if opts.repository != "chromeos":
+            raise RuntimeError(
+                "only chromeos supports automatic cherry-picking."
+            )
+
         new_state = do_cherrypick(
-            chroot_path=chroot_path,
+            chromeos_path=chromeos_path,
             llvm_dir=llvm_dir,
             repository=repository,
             interesting_shas=interesting_shas,
diff --git a/llvm_tools/nightly_revert_checker_test.py b/llvm_tools/nightly_revert_checker_test.py
index f1ee0bbc..47551129 100644
--- a/llvm_tools/nightly_revert_checker_test.py
+++ b/llvm_tools/nightly_revert_checker_test.py
@@ -144,7 +144,7 @@ class Test(unittest.TestCase):
             revert_checker.Revert("12345abcdef", "fedcba54321")
         ]
         nightly_revert_checker.do_cherrypick(
-            chroot_path="/path/to/chroot",
+            chromeos_path="/path/to/chroot",
             llvm_dir="/path/to/llvm",
             repository="repository_name",
             interesting_shas=[("12345abcdef", "fedcba54321")],