add support for main in git_cache

Change-Id: I32dcd812b41518bd8e2ca50eae1862611ff577de
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/tools/depot_tools/+/2528797
Reviewed-by: Edward Lesmes <ehmaldonado@chromium.org>
Commit-Queue: Anthony Polito <apolito@google.com>
diff --git a/git_cache.py b/git_cache.py
index 4ff3883..beee6fe 100755
--- a/git_cache.py
+++ b/git_cache.py
@@ -475,10 +475,10 @@
         self._fetch(self.mirror_path, verbose, depth, no_fetch_tags,
                     reset_fetch_config)
 
-  def update_bootstrap(self, prune=False, gc_aggressive=False):
+  def update_bootstrap(self, prune=False, gc_aggressive=False, branch='master'):
     # The folder is <git number>
     gen_number = subprocess.check_output(
-        [self.git_exe, 'number', 'master'],
+        [self.git_exe, 'number', branch],
         cwd=self.mirror_path).decode('utf-8', 'ignore').strip()
     gsutil = Gsutil(path=self.gsutil_exe, boto_path=None)
 
@@ -598,6 +598,8 @@
                     help='Run aggressive repacking of the repo.')
   parser.add_option('--prune', action='store_true',
                     help='Prune all other cached bundles of the same repo.')
+  parser.add_option('--branch', default='master',
+                    help='Branch to use for bootstrap. (Default \'master\')')
 
   populate_args = args[:]
   options, args = parser.parse_args(args)
@@ -612,7 +614,7 @@
   _, args2 = parser.parse_args(args)
   url = args2[0]
   mirror = Mirror(url)
-  mirror.update_bootstrap(options.prune, options.gc_aggressive)
+  mirror.update_bootstrap(options.prune, options.gc_aggressive, options.branch)
   return 0
 
 
diff --git a/tests/git_cache_test.py b/tests/git_cache_test.py
index 8cd95ae..83b1bb5 100755
--- a/tests/git_cache_test.py
+++ b/tests/git_cache_test.py
@@ -42,6 +42,12 @@
   def testParseFetchSpec(self):
     testData = [
         ([], []),
+        (['main'], [('+refs/heads/main:refs/heads/main',
+                       r'\+refs/heads/main:.*')]),
+        (['main/'], [('+refs/heads/main:refs/heads/main',
+                       r'\+refs/heads/main:.*')]),
+        (['+main'], [('+refs/heads/main:refs/heads/main',
+                       r'\+refs/heads/main:.*')]),
         (['master'], [('+refs/heads/master:refs/heads/master',
                        r'\+refs/heads/master:.*')]),
         (['master/'], [('+refs/heads/master:refs/heads/master',