pgo_tools: refactor chroot checking into a shared function

A new script can make use of "are we in the chroot" checking; might as
well centralize the logic now.

BUG=b:301994149
TEST=Unittests

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

4 files changed, 21 insertions, 6 deletions

diff --git a/pgo_tools/benchmark_pgo_profiles.py b/pgo_tools/benchmark_pgo_profiles.py
index 7b46b99e..73b99c7a 100755
--- a/pgo_tools/benchmark_pgo_profiles.py
+++ b/pgo_tools/benchmark_pgo_profiles.py
@@ -187,8 +187,7 @@ def main(argv: List[str]):
     )
     opts = parser.parse_args(argv)
 
-    if not Path("/etc/cros_chroot_version").exists():
-        sys.exit("Run me inside of the chroot.")
+    pgo_tools.exit_if_not_in_chroot()
 
     profiles = opts.profile
     validate_profiles(parser, profiles)
diff --git a/pgo_tools/create_chroot_and_generate_pgo_profile.py b/pgo_tools/create_chroot_and_generate_pgo_profile.py
index 657f34fa..ea812fcd 100755
--- a/pgo_tools/create_chroot_and_generate_pgo_profile.py
+++ b/pgo_tools/create_chroot_and_generate_pgo_profile.py
@@ -181,8 +181,7 @@ def main(argv: List[str]):
     )
     opts = parser.parse_args(argv)
 
-    if Path("/etc/cros_chroot_version").exists():
-        sys.exit("Do not run this inside of the chroot.")
+    pgo_tools.assert_not_in_chroot()
 
     repo_root = find_repo_root(Path(os.getcwd()))
     logging.info("Repo root is %s", repo_root)
diff --git a/pgo_tools/generate_pgo_profile.py b/pgo_tools/generate_pgo_profile.py
index a5fb7d45..5ff4221a 100755
--- a/pgo_tools/generate_pgo_profile.py
+++ b/pgo_tools/generate_pgo_profile.py
@@ -403,8 +403,7 @@ def main(argv: List[str]):
     )
     opts = parser.parse_args(argv)
 
-    if not Path("/etc/cros_chroot_version").exists():
-        sys.exit("Run me inside of the chroot.")
+    pgo_tools.exit_if_not_in_chroot()
 
     output = opts.output
 
diff --git a/pgo_tools/pgo_tools.py b/pgo_tools/pgo_tools.py
index 0106ae95..d34005b9 100644
--- a/pgo_tools/pgo_tools.py
+++ b/pgo_tools/pgo_tools.py
@@ -10,6 +10,7 @@ from pathlib import Path
 import re
 import shlex
 import subprocess
+import sys
 from typing import Any, Dict, IO, List, Optional, Union
 
 
@@ -119,3 +120,20 @@ def generate_quickpkg_restoration_command(quickpkg_path: Path) -> Command:
     package_ver = quickpkg_path.stem
     category = quickpkg_path.parent.name
     return ["sudo", "emerge", "--usepkgonly", f"={category}/{package_ver}"]
+
+
+def is_in_chroot() -> bool:
+    """Returns whether this script was invoked inside of the chroot."""
+    return Path("/etc/cros_chroot_version").exists()
+
+
+def exit_if_not_in_chroot():
+    """Calls sys.exit if this script was not run inside of the chroot."""
+    if not is_in_chroot():
+        sys.exit("Run me inside of the chroot.")
+
+
+def exit_if_in_chroot():
+    """Calls sys.exit if this script was run inside of the chroot."""
+    if is_in_chroot():
+        sys.exit("Run me outside of the chroot.")