Add -B/--bucket flag to git-cl try

git_cl.py

Index: git_cl.py
diff --git a/git_cl.py b/git_cl.py
index 174d46b3be6f0fa740e246577b7278f63dbbde5a..5b9054241375c96675a96707fc002d9647d91baf 100755
--- a/git_cl.py
+++ b/git_cl.py
@@ -76,6 +76,9 @@ REFS_THAT_ALIAS_TO_OTHER_REFS = {
 DEFAULT_LINT_REGEX = r"(.*\.cpp|.*\.cc|.*\.h)"
 DEFAULT_LINT_IGNORE_REGEX = r"$^"
 
+# Buildbucket master name prefix.
+MASTER_PREFIX = 'master.'
+
 # Shortcut since it quickly becomes redundant.
 Fore = colorama.Fore
 
@@ -275,10 +278,22 @@ def _prefix_master(master):
   (tryserver.chromium.linux) by stripping off the prefix 'master.'. This
   function does the conversion for buildbucket migration.
   """
-  prefix = 'master.'
-  if master.startswith(prefix):
+  if master.startswith(MASTER_PREFIX):
     return master
-  return '%s%s' % (prefix, master)
+  return '%s%s' % (MASTER_PREFIX, master)
+
+
+def _unprefix_master(bucket):
+  """Convert bucket name to shortened master name.
+
+  Buildbucket uses full master name(master.tryserver.chromium.linux) as bucket
+  name, while the developers always use shortened master name
+  (tryserver.chromium.linux) by stripping off the prefix 'master.'. This
+  function does the conversion for buildbucket migration.
+  """
+  if bucket.startswith(MASTER_PREFIX):
+    return bucket[len(MASTER_PREFIX):]
+  return bucket
 
 
 def _buildbucket_retry(operation_name, http, *args, **kwargs):
@@ -318,7 +333,7 @@ def _buildbucket_retry(operation_name, http, *args, **kwargs):
   assert False, 'unreachable'
 
 
-def _trigger_try_jobs(auth_config, changelist, masters, options,
+def _trigger_try_jobs(auth_config, changelist, buckets, options,
                       category='git_cl_try', patchset=None):
   assert changelist.GetIssue(), 'CL must be uploaded first'
   codereview_url = changelist.GetCodereviewServer()
@@ -350,9 +365,9 @@ def _trigger_try_jobs(auth_config, changelist, masters, options,
   batch_req_body = {'builds': []}
   print_text = []
   print_text.append('Tried jobs on:')
-  for master, builders_and_tests in sorted(masters.iteritems()):
-    print_text.append('Master: %s' % master)
-    bucket = _prefix_master(master)
+  for bucket, builders_and_tests in sorted(buckets.iteritems()):
+    print_text.append('Bucket: %s' % bucket)
+    master = _unprefix_master(bucket)
     for builder, tests in sorted(builders_and_tests.iteritems()):
       print_text.append('  %s: %s' % (builder, tests))
       parameters = {
@@ -4705,9 +4720,11 @@ def CMDtry(parser, args):
             'the try server waterfall for the builders name and the tests '
             'available.'))
   group.add_option(
+      '-B', '--bucket', default='',
+      help=('Buildbucket bucket to send the try requests.'))
+  group.add_option(
       '-m', '--master', default='',
       help=('Specify a try master where to run the tries.'))
-  # TODO(tandrii,nodir): add -B --bucket flag.
   group.add_option(
       '-r', '--revision',
       help='Revision to use for the try job; default: the revision will '
@@ -4765,13 +4782,36 @@ def CMDtry(parser, args):
   if not options.name:
     options.name = cl.GetBranch()
 
-  if options.bot and not options.master:
+  if options.bucket and options.master:
+    parser.error('Only one of --bucket and --master may be used.')
+
+  if options.bot and not options.master and not options.bucket:
     options.master, err_msg = GetBuilderMaster(options.bot)
     if err_msg:
       parser.error('Tryserver master cannot be found because: %s\n'
                    'Please manually specify the tryserver master'
                    ', e.g. "-m tryserver.chromium.linux".' % err_msg)
 
+  def GetBotsMap(bots):
+    builders_and_tests = {}
+
+    # TODO(machenbach): The old style command-line options don't support
+    # multiple try masters yet.
+    old_style = filter(lambda x: isinstance(x, basestring), bots)
+    new_style = filter(lambda x: isinstance(x, tuple), bots)
+
+    for bot in old_style:
+      if ':' in bot:
+        parser.error('Specifying testfilter is no longer supported')
+      elif ',' in bot:
+        parser.error('Specify one bot per --bot flag')
+      else:
+        builders_and_tests.setdefault(bot, [])
+
+    for bot, tests in new_style:
+      builders_and_tests.setdefault(bot, []).extend(tests)
+    return builders_and_tests
+
   def GetMasterMap():
     # Process --bot.
     if not options.bot:
@@ -4802,53 +4842,45 @@ def CMDtry(parser, args):
     if not options.bot:
       return {}
 
-    builders_and_tests = {}
-    # TODO(machenbach): The old style command-line options don't support
-    # multiple try masters yet.
-    old_style = filter(lambda x: isinstance(x, basestring), options.bot)
-    new_style = filter(lambda x: isinstance(x, tuple), options.bot)
-
-    for bot in old_style:
-      if ':' in bot:
-        parser.error('Specifying testfilter is no longer supported')
-      elif ',' in bot:
-        parser.error('Specify one bot per --bot flag')
-      else:
-        builders_and_tests.setdefault(bot, [])
-
-    for bot, tests in new_style:
-      builders_and_tests.setdefault(bot, []).extend(tests)
+    builders_and_tests = GetBotsMap(options.bot)
 
     # Return a master map with one master to be backwards compatible. The
     # master name defaults to an empty string, which will cause the master
     # not to be set on rietveld (deprecated).
-    return {options.master: builders_and_tests}
-
-  masters = GetMasterMap()
-  if not masters:
-    # Default to triggering Dry Run (see http://crbug.com/625697).
-    if options.verbose:
-      print('git cl try with no bots now defaults to CQ Dry Run.')
-    try:
-      cl.SetCQState(_CQState.DRY_RUN)
-      print('scheduled CQ Dry Run on %s' % cl.GetIssueURL())
-      return 0
-    except KeyboardInterrupt:
-      raise
-    except:
-      print('WARNING: failed to trigger CQ Dry Run.\n'
-            'Either:\n'
-            ' * your project has no CQ\n'
-            ' * you don\'t have permission to trigger Dry Run\n'
-            ' * bug in this code (see stack trace below).\n'
-            'Consider specifying which bots to trigger manually '
-            'or asking your project owners for permissions '
-            'or contacting Chrome Infrastructure team at '
-            'https://www.chromium.org/infra\n\n')
-      # Still raise exception so that stack trace is printed.
-      raise
-
-  for builders in masters.itervalues():
+    bucket = ''
+    if options.master:
+      # Add the "master." prefix to the master name to obtain the bucket name.
+      bucket = _prefix_master(options.master)
+    return {bucket: builders_and_tests}
+
+  if options.bucket:
+    buckets = {options.bucket: GetBotsMap(options.bot or [])}
+  else:
+    buckets = GetMasterMap()
+    if not buckets:
+      # Default to triggering Dry Run (see http://crbug.com/625697).
+      if options.verbose:
+        print('git cl try with no bots now defaults to CQ Dry Run.')
+      try:
+        cl.SetCQState(_CQState.DRY_RUN)
+        print('scheduled CQ Dry Run on %s' % cl.GetIssueURL())
+        return 0
+      except KeyboardInterrupt:
+        raise
+      except:
+        print('WARNING: failed to trigger CQ Dry Run.\n'
+              'Either:\n'
+              ' * your project has no CQ\n'
+              ' * you don\'t have permission to trigger Dry Run\n'
+              ' * bug in this code (see stack trace below).\n'
+              'Consider specifying which bots to trigger manually '
+              'or asking your project owners for permissions '
+              'or contacting Chrome Infrastructure team at '
+              'https://www.chromium.org/infra\n\n')
+        # Still raise exception so that stack trace is printed.
+        raise
+
+  for builders in buckets.itervalues():
     if any('triggered' in b for b in builders):
       print('ERROR You are trying to send a job to a triggered bot. This type '
             'of bot requires an initial job from a parent (usually a builder). '
@@ -4865,8 +4897,8 @@ def CMDtry(parser, args):
           'codereview, continuing to use patchset %s.\n' %
           (patchset, cl.GetPatchset(), patchset))
   try:
-    _trigger_try_jobs(auth_config, cl, masters, options, 'git_cl_try',
-                     patchset)
+    _trigger_try_jobs(auth_config, cl, buckets, options, 'git_cl_try',
+                      patchset)
   except BuildbucketResponseException as ex:
     print('ERROR: %s' % ex)
     return 1