toolchain_util: stop hardcoding master branch

Get the details from the manifest itself.

BUG=chromium:1126855
TEST=CQ passes

Change-Id: I2dccff3a72d579e26cc423b6db1a6601f5c93221
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/2595353
Reviewed-by: Tiancong Wang <tcwang@google.com>
Reviewed-by: Mike Frysinger <vapier@chromium.org>
Commit-Queue: Tiancong Wang <tcwang@google.com>
Tested-by: Tiancong Wang <tcwang@google.com>
diff --git a/lib/toolchain_util.py b/lib/toolchain_util.py
index 52fa796..a5aaf04 100644
--- a/lib/toolchain_util.py
+++ b/lib/toolchain_util.py
@@ -79,9 +79,8 @@
 BZ2_COMPRESSION_SUFFIX = '.bz2'
 XZ_COMPRESSION_SUFFIX = '.xz'
 KERNEL_AFDO_COMPRESSION_SUFFIX = '.gcov.xz'
-TOOLCHAIN_UTILS_PATH = '/mnt/host/source/src/third_party/toolchain-utils/'
-TOOLCHAIN_UTILS_REPO = \
-    'https://chromium.googlesource.com/chromiumos/third_party/toolchain-utils'
+TOOLCHAIN_UTILS_PATH = os.path.join(
+    constants.SOURCE_ROOT, 'src/third_party/toolchain-utils')
 AFDO_PROFILE_PATH_IN_CHROMIUM = 'src/chromeos/profiles/%s.afdo.newest.txt'
 MERGED_AFDO_NAME = 'chromeos-chrome-amd64-%s'
 
@@ -3018,6 +3017,9 @@
     if the uploaded artifact is the same as in metadata file (no new
     AFDO artifact uploaded). This should be caught earlier before uploading.
   """
+  branch = git.GetTrackingBranch(
+      TOOLCHAIN_UTILS_PATH, for_checkout=False, for_push=True)
+
   # Perform a git pull first to sync the checkout containing metadata,
   # in case there are some updates during the builder was running.
   # Possible race conditions are:
@@ -3028,9 +3030,7 @@
   # we might update an entry with an older profile. The checking logic in the
   # function should guarantee we always update to a newer artifact, otherwise
   # the builder fails.
-  git.RunGit(
-      TOOLCHAIN_UTILS_PATH, ['pull', TOOLCHAIN_UTILS_REPO, 'refs/heads/master'],
-      print_cmd=True)
+  git.RunGit(TOOLCHAIN_UTILS_PATH, ['pull', branch.remote], print_cmd=True)
 
   afdo_versions = json.loads(osutils.ReadFile(json_file))
   if title:
@@ -3088,7 +3088,7 @@
   git.GitPush(
       TOOLCHAIN_UTILS_PATH,
       'HEAD',
-      git.RemoteRef(TOOLCHAIN_UTILS_REPO, 'refs/for/master%submit'),
+      git.RemoteRef(branch.remote, f'{branch.ref}%submit'),
       print_cmd=True)
 
 
diff --git a/lib/toolchain_util_unittest.py b/lib/toolchain_util_unittest.py
index 6c63175..2873ac9 100644
--- a/lib/toolchain_util_unittest.py
+++ b/lib/toolchain_util_unittest.py
@@ -2303,6 +2303,8 @@
 
     with open(self.json_file, 'w') as f:
       json.dump(self.afdo_versions, f)
+    self.PatchObject(git, 'GetTrackingBranch',
+                     return_value=git.RemoteRef('origin', 'main'))
     GitStatus = collections.namedtuple('GitStatus', ['output'])
     self.mock_git = self.PatchObject(
         git, 'RunGit', return_value=GitStatus(output='non-empty'))
@@ -2337,12 +2339,7 @@
                                               self.afdo_sorted_by_freshness[2])
     calls = [
         mock.call(
-            self.tempdir, [
-                'pull',
-                toolchain_util.TOOLCHAIN_UTILS_REPO,
-                'refs/heads/master',
-            ],
-            print_cmd=True),
+            self.tempdir, ['pull', 'origin'], print_cmd=True),
         mock.call(
             self.tempdir, ['status', '--porcelain', '-uno'],
             capture_output=True,
@@ -2352,8 +2349,8 @@
             self.tempdir, ['commit', '-a', '-m', message], print_cmd=True),
         mock.call(
             self.tempdir, [
-                'push', toolchain_util.TOOLCHAIN_UTILS_REPO,
-                'HEAD:refs/for/master%submit'
+                'push', 'origin',
+                'HEAD:main%submit'
             ],
             capture_output=True,
             print_cmd=True)