Refactor getting branch override logic

scale_build/packages/git.py

@@ -2,8 +2,10 @@ import logging
 import os
 import shutil
 
+from scale_build.config import BRANCH_OVERRIDES, TRY_BRANCH_OVERRIDE
+from scale_build.exceptions import CallError
 from scale_build.utils.git_utils import (
-    create_branch, retrieve_git_remote_and_sha, retrieve_git_branch, update_git_manifest
+    branch_exists_in_repository, create_branch, retrieve_git_remote_and_sha, retrieve_git_branch, update_git_manifest
 )
 from scale_build.utils.logger import LoggingContext
 from scale_build.utils.paths import BRANCH_OUT_LOG_FILENAME, GIT_LOG_PATH
@@ -53,3 +55,30 @@ class GitPackageMixin:
         if not self.exists:
             return None
         return retrieve_git_branch(self.source_path)
+
+    def get_branch_override(self):
+        gh_override = BRANCH_OVERRIDES.get(self.name)
+
+        # TRY_BRANCH_OVERRIDE is a special use-case. It allows setting a branch name to be used
+        # during the checkout phase, only if it exists on the remote.
+        #
+        # This is useful for PR builds and testing where you want to use defaults for most repos
+        # but need to test building of a series of repos with the same experimental branch
+        #
+        if TRY_BRANCH_OVERRIDE:
+            retries = 3
+            while retries:
+                try:
+                    if branch_exists_in_repository(self.origin, TRY_BRANCH_OVERRIDE):
+                        gh_override = TRY_BRANCH_OVERRIDE
+                except CallError:
+                    retries -= 1
+                    logger.debug(
+                        'Failed to determine if %r branch exists for %r. Trying again', TRY_BRANCH_OVERRIDE, self.origin
+                    )
+                    if not retries:
+                        logger.debug('Unable to determine if %r branch exists in 3 attempts.', TRY_BRANCH_OVERRIDE)
+                else:
+                    break
+
+        return gh_override