Add --use-swarming-go to the custom trigger scripts

Design doc: http://go/roll-go-swarming-to-trigger-scripts

Bug: 1127205
Change-Id: I8ca8e2bb381f340fc23b91339cdb56374f8eead8
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2500976Reviewed-by: Takuto Ikuta <tikuta@chromium.org>
Reviewed-by: Wenbin Zhang <wenbinzhang@google.com>
Commit-Queue: Ye Kuang <yekuang@google.com>
Cr-Commit-Position: refs/heads/master@{#821619}

testing/trigger_scripts/base_test_triggerer.py

@@ -304,7 +304,10 @@ class BaseTestTriggerer(object):
                                         suffix='.json')
         args_to_pass = self.modify_args(filtered_remaining_args, bot_index,
                                         shard_index, args.shards, json_temp)
-        ret = self.run_swarming(args_to_pass, verbose)
+        if args.use_swarming_go:
+          ret = self.run_swarming_go(args_to_pass, verbose)
+        else:
+          ret = self.run_swarming(args_to_pass, verbose)
         if ret:
           sys.stderr.write('Failed to trigger a task, aborting\n')
           return ret
@@ -345,5 +348,12 @@ class BaseTestTriggerer(object):
     parser.add_argument('--shard-index', type=int, default=None,
                         help='Which shard to trigger. Duplicated from the '
                              '`swarming.py trigger` command.')
+    BaseTestTriggerer.add_use_swarming_go_arg(parser)
     return parser
 
+  @staticmethod
+  def add_use_swarming_go_arg(parser):
+    parser.add_argument('--use-swarming-go', type=bool, default=False,
+                        action='store_true',
+                        help='Uses swarming Go CLI to trigger tasks.')
+

testing/trigger_scripts/chromeos_device_trigger.py

@@ -75,6 +75,7 @@ def parse_args(triggerer):
     additional_args = triggerer.modify_args(
         additional_args, 0, args.shard_index, args.shards, args.dump_json)
   else:
+    base_test_triggerer.BaseTestTriggerer.add_use_swarming_go_arg(parser)
     args, additional_args = parser.parse_known_args()
 
   if additional_args[0] != 'trigger':
@@ -128,6 +129,8 @@ def main():
   ])
   new_args += additional_args[1:]
 
+  if args.use_swarming_go:
+    return triggerer.run_swarming_go(new_args, True)
   return triggerer.run_swarming(new_args, True)
 
 

testing/trigger_scripts/perf_device_trigger_unittest.py

@@ -16,6 +16,7 @@ class Args(object):
     self.dump_json = ''
     self.multiple_trigger_configs = None
     self.multiple_dimension_script_verbose = False
+    self.use_swarming_go = False
 
 
 class FakeTriggerer(perf_device_trigger.PerfDeviceTriggerer):
@@ -24,6 +25,7 @@ class FakeTriggerer(perf_device_trigger.PerfDeviceTriggerer):
     self._swarming_runs = []
     self._files = files
     self._temp_file_id = 0
+    self._triggered_with_swarming_go = 0
     super(FakeTriggerer, self).__init__(args, swarming_args)
 
 
@@ -51,14 +53,21 @@ class FakeTriggerer(perf_device_trigger.PerfDeviceTriggerer):
     del verbose #unused
     self._swarming_runs.append(args)
 
+  def run_swarming_go(self, args, verbose):
+    self._triggered_with_swarming_go += 1
+    self.run_swarming(args, verbose)
 
 class UnitTest(unittest.TestCase):
-  def setup_and_trigger(
-      self, previous_task_assignment_map, alive_bots, dead_bots):
+  def setup_and_trigger(self,
+                        previous_task_assignment_map,
+                        alive_bots,
+                        dead_bots,
+                        use_swarming_go=False):
     args = Args()
     args.shards = len(previous_task_assignment_map)
     args.dump_json = 'output.json'
     args.multiple_dimension_script_verbose = True
+    args.use_swarming_go = use_swarming_go
     swarming_args = [
         'trigger',
         '--swarming',
@@ -291,6 +300,15 @@ class UnitTest(unittest.TestCase):
     self.assertEquals(set(expected_task_assignment.values()),
         {'build3', 'build4', 'build5', 'build7'})
 
+  def test_use_swarming_go_to_trigger(self):
+    triggerer = self.setup_and_trigger(
+        previous_task_assignment_map={0: 'build1'},
+        alive_bots=['build1'],
+        dead_bots=[],
+        use_swarming_go=True)
+
+    self.assertEquals(triggerer._triggered_with_swarming_go, 1)
+
 
 if __name__ == '__main__':
   unittest.main()