[Findit] Refactoring trigger swarming task pipelines

appengine/findit/waterfall/test/trigger_base_swarming_task_result_pipeline_test.py

 # Copyright 2016 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 import time
 
 from model import analysis_status
+from model.flake.flake_swarming_task import FlakeSwarmingTask
 from model.wf_swarming_task import WfSwarmingTask
 from waterfall import swarming_util
-from waterfall import trigger_swarming_task_pipeline
 from waterfall import waterfall_config
 from waterfall.swarming_task_request import SwarmingTaskRequest
 from waterfall.test import wf_testcase
+from waterfall.trigger_base_swarming_task_pipeline import (
+    TriggerBaseSwarmingTaskPipeline)
+from waterfall.trigger_flake_swarming_task_pipeline import (
+    TriggerFlakeSwarmingTaskPipeline)
 from waterfall.trigger_swarming_task_pipeline import TriggerSwarmingTaskPipeline
 
 
-class TriggerSwarmingTaskPipelineTest(wf_testcase.WaterfallTestCase):
+class TriggerBaseSwarmingTaskPipelineTest(wf_testcase.WaterfallTestCase):
 
   def testNoNewSwarmingTaskIsNeeded(self):
     master_name = 'm'
     builder_name = 'b'
     build_number = 1
     step_name = 's'
     tests = ['a.b']
 
-    def MockedListSwarmingTasksDataByTags(*_):
-      raise Exception('Fail.')  # pragma: no cover. Should never be reached.
-    self.mock(swarming_util, 'ListSwarmingTasksDataByTags',
-              MockedListSwarmingTasksDataByTags)
     swarming_task = WfSwarmingTask.Create(
         master_name, builder_name, build_number, step_name)
     swarming_task.status = analysis_status.RUNNING
     swarming_task.task_id = 'task_id'
     swarming_task.put()
 
     pipeline = TriggerSwarmingTaskPipeline()
     task_id = pipeline.run(
         master_name, builder_name, build_number, step_name, tests)
     self.assertEqual('task_id', task_id)
 
   def testWaitingForTheTaskId(self):
     master_name = 'm'
     builder_name = 'b'
     build_number = 1
     step_name = 's'
     tests = ['a.b']
 
-    swarming_task = WfSwarmingTask.Create(
-        master_name, builder_name, build_number, step_name)
+    swarming_task = FlakeSwarmingTask.Create(
+        master_name, builder_name, build_number, step_name, tests[0])
     swarming_task.status = analysis_status.PENDING
     swarming_task.put()
 
     def MockedSleep(*_):
-      swarming_task = WfSwarmingTask.Get(
-          master_name, builder_name, build_number, step_name)
+      swarming_task = FlakeSwarmingTask.Get(
+          master_name, builder_name, build_number, step_name, tests[0])
       self.assertEqual(analysis_status.PENDING, swarming_task.status)
       swarming_task.status = analysis_status.RUNNING
       swarming_task.task_id = 'task_id'
       swarming_task.put()
     self.mock(time, 'sleep', MockedSleep)
 
-    pipeline = TriggerSwarmingTaskPipeline()
+    pipeline = TriggerFlakeSwarmingTaskPipeline()
     task_id = pipeline.run(
         master_name, builder_name, build_number, step_name, tests)
     self.assertEqual('task_id', task_id)
 
   def testTriggerANewSwarmingTask(self):
     def MockedDownloadSwarmingTaskData(*_):
       return [{'task_id': '1'}, {'task_id': '2'}]
     self.mock(swarming_util, 'ListSwarmingTasksDataByTags',
               MockedDownloadSwarmingTaskData)
 
@@ skipping 26 matching lines @@
               MockedGetSwarmingTaskRequest)
 
     new_request_json = {}
     def MockedTriggerSwarmingTask(new_request, *_):
       new_request_json.update(new_request.Serialize())
       return 'new_task_id'
     self.mock(swarming_util, 'TriggerSwarmingTask', MockedTriggerSwarmingTask)
 
     def MockedGetSwarmingTaskName(*_):
       return 'new_task_name'
-    self.mock(trigger_swarming_task_pipeline, '_GetSwarmingTaskName',
+    self.mock(TriggerBaseSwarmingTaskPipeline, '_GetSwarmingTaskName',
               MockedGetSwarmingTaskName)
 
     master_name = 'm'
     builder_name = 'b'
     build_number = 234
     step_name = 'a_tests on platform'
     tests = ['a.b', 'a.c']
 
     expected_new_request_json = {
         'expiration_secs': 3600,
@@ skipping 32 matching lines @@
     self.assertEqual(expected_new_request_json, new_request_json)
 
     swarming_task = WfSwarmingTask.Get(
         master_name, builder_name, build_number, step_name)
     self.assertIsNotNone(swarming_task)
     self.assertEqual('new_task_id', swarming_task.task_id)
     self.assertEqual(tests, swarming_task.parameters['tests'])
     self.assertEqual(
         waterfall_config.GetSwarmingSettings()['iterations_to_rerun'],
         swarming_task.parameters['iterations_to_rerun'])