[Findit] Refactoring trigger swarming task pipelines

appengine/findit/waterfall/trigger_swarming_task_pipeline.py

Index: appengine/findit/waterfall/trigger_swarming_task_pipeline.py
diff --git a/appengine/findit/waterfall/trigger_swarming_task_pipeline.py b/appengine/findit/waterfall/trigger_swarming_task_pipeline.py
index 0c683b56a08c01978461bf16e12d97f1e3ba86e8..f847dfed424ae8828dfd659dbbed6a3db5f08fe8 100644
--- a/appengine/findit/waterfall/trigger_swarming_task_pipeline.py
+++ b/appengine/findit/waterfall/trigger_swarming_task_pipeline.py
@@ -2,166 +2,33 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-import copy
-import logging
-import time
-
-from google.appengine.ext import ndb
-
-from common.http_client_appengine import HttpClientAppengine as HttpClient
-from common.pipeline_wrapper import BasePipeline
-from lib import time_util
-from model import analysis_status
 from model.wf_swarming_task import WfSwarmingTask
-from waterfall import swarming_util
 from waterfall import waterfall_config
+from waterfall.trigger_base_swarming_task_pipeline import (
+    TriggerBaseSwarmingTaskPipeline)
 
 
-def _GetSwarmingTaskName(ref_task_id):  # pragma: no cover.
-  """Returns a unique task name.
-
-  Have this separate function in order to mock for testing purpose.
-  """
-  return 'findit/deflake/ref_task_id/%s/%s' % (
-      ref_task_id, time_util.GetUTCNow().strftime('%Y-%m-%d %H:%M:%S %f'))
-
-
-def _CreateNewSwarmingTaskRequest(
-    ref_task_id, ref_request, master_name, builder_name, build_number,
-    step_name, tests, iterations):
-  """Returns a SwarmingTaskRequest instance to run the given tests only."""
-  # Make a copy of the referred request and drop or overwrite some fields.
-  new_request = copy.deepcopy(ref_request)
-  new_request.name = _GetSwarmingTaskName(ref_task_id)
-  new_request.parent_task_id = ''
-  new_request.user = ''
-
-  # To force a fresh re-run and ignore cached result of any equivalent run.
-  new_request.idempotent = False
-
-  # Set the gtest_filter to run the given tests only.
-  new_request.extra_args.append('--gtest_repeat=%s' % iterations)
-  new_request.extra_args.append('--test-launcher-retry-limit=0')
-  new_request.extra_args = [
-      a for a in new_request.extra_args if not a.startswith('--gtest_filter')
-  ]
-  new_request.extra_args.append('--gtest_filter=%s' % ':'.join(tests))
-
-  # Remove the env setting for sharding.
-  sharding_settings = ['GTEST_SHARD_INDEX', 'GTEST_TOTAL_SHARDS']
-  new_request.env = [
-      e for e in new_request.env if e['key'] not in sharding_settings
-  ]
-
-  # Reset tags for searching and monitoring.
-  ref_name = swarming_util.GetTagValue(ref_request.tags, 'name')
-  new_request.tags = []
-  new_request.tags.append('purpose:deflake')
-  new_request.tags.append('ref_master:%s' % master_name)
-  new_request.tags.append('ref_buildername:%s' % builder_name)
-  new_request.tags.append('ref_buildnumber:%s' % build_number)
-  new_request.tags.append('ref_stepname:%s' % step_name)
-  new_request.tags.append('ref_task_id:%s' % ref_task_id)
-  new_request.tags.append('ref_name:%s' % ref_name)
-
-  return new_request
-
-
-@ndb.transactional
-def _NeedANewSwarmingTask(master_name, builder_name, build_number, step_name):
-  swarming_task = WfSwarmingTask.Get(
-      master_name, builder_name, build_number, step_name)
-
-  if not swarming_task:
-    swarming_task = WfSwarmingTask.Create(
-        master_name, builder_name, build_number, step_name)
-    swarming_task.status = analysis_status.PENDING
-    swarming_task.put()
-    return True
-  else:
-    # TODO(http://crbug.com/585676): Rerun the Swarming task if it runs into
-    # unexpected infra errors.
-    return False
-
-
-def _GetSwarmingTaskId(master_name, builder_name, build_number, step_name):
-  deadline = time.time() + 5 * 60  # Wait for 5 minutes.
-  while time.time() < deadline:
-    swarming_task = WfSwarmingTask.Get(
-        master_name, builder_name, build_number, step_name)
-
-    if not swarming_task:  # pragma: no cover. Pipeline will retry.
-      raise Exception('Swarming task was deleted unexpectedly!!!')
-
-    if swarming_task.task_id:
-      return swarming_task.task_id
-
-    time.sleep(10)  # Wait for the existing pipeline to start the Swarming task.
-
-  raise Exception('Time out!')  # pragma: no cover. Pipeline will retry.
-
-
-class TriggerSwarmingTaskPipeline(BasePipeline):
+class TriggerSwarmingTaskPipeline(TriggerBaseSwarmingTaskPipeline):
   """A pipeline to trigger a Swarming task to re-run selected tests of a step.
 
   This pipeline only supports test steps that run on Swarming and support the
   gtest filter.
   """
 
-  # Arguments number differs from overridden method - pylint: disable=W0221
-  def run(self, master_name, builder_name, build_number, step_name, tests):
-    """Triggers a new Swarming task to run the given tests.
-
-    Args:
-      master_name (str): The master name.
-      builder_name (str): The builder name.
-      build_number (str): The build number.
-      step_name (str): The failed test step name.
-      tests (list): A list of test cases, eg: ['suite1.test1', 'suite2.test2'].
-
-    Returns:
-      task_id (str): The new Swarming task that re-run the given tests.
-    """
-    # Check if a new Swarming Task is really needed.
-    if not _NeedANewSwarmingTask(
-        master_name, builder_name, build_number, step_name):
-      return _GetSwarmingTaskId(
-          master_name, builder_name, build_number, step_name)
+  def _GetArgs(self, master_name, builder_name, build_number, step_name, _):
+    return master_name, builder_name, build_number, step_name
 
-    assert tests
-
-    http_client = HttpClient()
-
-    # 0. Retrieve existing Swarming task ids for the given step.
-    swarming_task_items = swarming_util.ListSwarmingTasksDataByTags(
-        master_name, builder_name, build_number, http_client,
-        {'stepname': step_name})
-    assert len(swarming_task_items) > 0, 'No Swarming task was run.'
-    ref_task_id = swarming_task_items[0]['task_id']
-
-    # 1. Retrieve Swarming task parameters from a given Swarming task id.
-    ref_request = swarming_util.GetSwarmingTaskRequest(
-        ref_task_id, http_client)
-
-    # 2. Update/Overwrite parameters for the re-run.
-    iterations_to_rerun = waterfall_config.GetSwarmingSettings().get(
-        'iterations_to_rerun')
-    new_request = _CreateNewSwarmingTaskRequest(
-        ref_task_id, ref_request, master_name, builder_name, build_number,
-        step_name, tests, iterations_to_rerun)
-
-    # 3. Trigger a new Swarming task to re-run the failed tests.
-    task_id = swarming_util.TriggerSwarmingTask(new_request, http_client)
+  # Arguments number differs from overridden method - pylint: disable=W0221
+  def _GetSwarmingTask(
+      self, master_name, builder_name, build_number, step_name):
+    return WfSwarmingTask.Get(
+        master_name, builder_name, build_number, step_name)
 
-    # Save the task id.
-    swarming_task = WfSwarmingTask.Get(
+  # Arguments number differs from overridden method - pylint: disable=W0221
+  def _CreateSwarmingTask(
+      self, master_name, builder_name, build_number, step_name):
+    return WfSwarmingTask.Create(
         master_name, builder_name, build_number, step_name)
-    swarming_task.task_id = task_id
-    swarming_task.parameters['tests'] = tests
-    swarming_task.parameters['iterations_to_rerun'] = iterations_to_rerun
-    swarming_task.parameters['ref_name'] = swarming_util.GetTagValue(
-        new_request.tags, 'ref_name')
-    swarming_task.put()
 
-    logging.info('A Swarming task was triggered:%s', task_id)
-    return task_id
+  def _GetIterationsToRerun(self):
+    return waterfall_config.GetSwarmingSettings().get('iterations_to_rerun')