From 1678a364a5fd10cbac81ff541b9315ca9d9b6e13 Mon Sep 17 00:00:00 2001
From: Kenneth Russell <kbr@chromium.org>
Date: Tue, 7 Feb 2017 15:21:24 -0800
Subject: [PATCH] Add presubmit_support.EnsureCQIncludeTrybotsAreAdded.

The PostUploadHooks in the Chromium repository which add
CQ_INCLUDE_TRYBOTS entries to the issue description will be rewritten
in terms of this primitive, which will compose properly if multiple
sub-directories attempt to modify it.

BUG=688765

Change-Id: Icf72edb872f29af1e082038e96bc547504edfd07
Reviewed-on: https://chromium-review.googlesource.com/438925
Reviewed-by: Marc-Antoine Ruel <maruel@chromium.org>
Reviewed-by: Andrii Shyshkalov <tandrii@chromium.org>
Commit-Queue: Kenneth Russell <kbr@chromium.org>
---
 presubmit_support.py        | 36 +++++++++++++++++++++++++++
 tests/presubmit_unittest.py | 49 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/presubmit_support.py b/presubmit_support.py
index 1f97d7f1c..d71ebd26e 100755
--- a/presubmit_support.py
+++ b/presubmit_support.py
@@ -1096,6 +1096,42 @@ def DoPostUploadExecuter(change,
   return results
 
 
+# This helper function should be used by any PostUploadHook wishing to
+# add entries to the CQ_INCLUDE_TRYBOTS line. It returns the results
+# that should be returned from the PostUploadHook.
+def EnsureCQIncludeTrybotsAreAdded(cl, bots_to_include, message, output_api):
+  """Helper to be used by any PostUploadHook wishing to add CQ_INCLUDE_TRYBOTS.
+
+  Merges the bots_to_include into the current CQ_INCLUDE_TRYBOTS list,
+  keeping it alphabetically sorted. Returns the results that should be
+  returned from the PostUploadHook.
+
+  Args:
+    cl: The git_cl.Changelist object.
+    bots_to_include: A list of strings of bots to include, in the form
+      "master:slave".
+    message: A message to be printed in the case that
+      CQ_INCLUDE_TRYBOTS was updated.
+    output_api: An OutputApi instance used to display messages.
+  """
+  rietveld_obj = cl.RpcServer()
+  issue = cl.issue
+  description = rietveld_obj.get_description(issue)
+  all_bots = []
+  include_re = re.compile(r'^CQ_INCLUDE_TRYBOTS=(.*)', re.M | re.I)
+  m = include_re.search(description)
+  if m:
+    all_bots = [i.strip() for i in m.group(1).split(';') if i.strip()]
+  if not (set(all_bots) - set(bots_to_include)):
+    return []
+  # Sort the bots to keep them in some consistent order -- not required.
+  all_bots = sorted(set(all_bots) | set(bots_to_include))
+  new_include_trybots = 'CQ_INCLUDE_TRYBOTS=%s' % ';'.join(all_bots)
+  new_description = include_re.sub(new_include_trybots, description)
+  rietveld_obj.update_description(issue, new_description)
+  return [output_api.PresubmitNotifyResult(message)]
+
+
 class PresubmitExecuter(object):
   def __init__(self, change, committing, rietveld_obj, verbose,
                gerrit_obj=None, dry_run=None):
diff --git a/tests/presubmit_unittest.py b/tests/presubmit_unittest.py
index 78319e609..0e7cc8525 100755
--- a/tests/presubmit_unittest.py
+++ b/tests/presubmit_unittest.py
@@ -168,7 +168,7 @@ class PresubmitUnittest(PresubmitTestsBase):
       'subprocess', 'sys', 'tempfile', 'time', 'traceback', 'types', 'unittest',
       'urllib2', 'warn', 'multiprocessing', 'DoGetTryMasters',
       'GetTryMastersExecuter', 'itertools', 'urlparse', 'gerrit_util',
-      'GerritAccessor',
+      'GerritAccessor', 'EnsureCQIncludeTrybotsAreAdded',
     ]
     # If this test fails, you should add the relevant test.
     self.compareMembers(presubmit, members)
@@ -955,6 +955,53 @@ def CheckChangeOnCommit(input_api, output_api):
     except SystemExit, e:
       self.assertEquals(2, e.code)
 
+  def testEnsureCQIncludeTrybotsAreAdded(self):
+    # Deliberately has a space at the end to exercise space-stripping code.
+    cl_text = """A change to GPU-related code.
+
+CQ_INCLUDE_TRYBOTS=master.tryserver.blink:linux_trusty_blink_rel;master.tryserver.chromium.win:win_optional_gpu_tests_rel 
+"""
+
+    updated_cl_text = """A change to GPU-related code.
+
+CQ_INCLUDE_TRYBOTS=master.tryserver.blink:linux_trusty_blink_rel;master.tryserver.chromium.linux:linux_optional_gpu_tests_rel;master.tryserver.chromium.mac:mac_optional_gpu_tests_rel;master.tryserver.chromium.win:win_optional_gpu_tests_rel
+"""
+
+    class FakeIssue(object):
+      def __init__(self, description):
+        self.description = description
+
+    class FakeRietveld(object):
+      def get_description(self, issue):
+        return issue.description
+
+      def update_description(self, issue, new_description):
+        issue.description = new_description
+
+    class FakeCL(object):
+      def __init__(self, description):
+        self.issue = FakeIssue(description)
+      def RpcServer(self):
+        return FakeRietveld()
+
+    class FakeOutputAPI(object):
+      def PresubmitNotifyResult(self, message):
+        return message
+
+    cl = FakeCL(cl_text)
+    message = 'Automatically added optional GPU tests to run on CQ.'
+    results = presubmit.EnsureCQIncludeTrybotsAreAdded(
+      cl,
+      [
+        'master.tryserver.chromium.linux:linux_optional_gpu_tests_rel',
+        'master.tryserver.chromium.mac:mac_optional_gpu_tests_rel',
+        'master.tryserver.chromium.win:win_optional_gpu_tests_rel'
+      ],
+      message,
+      FakeOutputAPI())
+    self.assertEqual(updated_cl_text, cl.issue.description)
+    self.assertEqual([message], results)
+
 
 class InputApiUnittest(PresubmitTestsBase):
   """Tests presubmit.InputApi."""