[Findit] Refactoring trigger swarming task pipelines

appengine/findit/waterfall/test/trigger_base_swarming_task_result_pipeline_test.py

Index: appengine/findit/waterfall/test/trigger_base_swarming_task_result_pipeline_test.py
diff --git a/appengine/findit/waterfall/test/trigger_swarming_task_pipeline_test.py b/appengine/findit/waterfall/test/trigger_base_swarming_task_result_pipeline_test.py
similarity index 88%
copy from appengine/findit/waterfall/test/trigger_swarming_task_pipeline_test.py
copy to appengine/findit/waterfall/test/trigger_base_swarming_task_result_pipeline_test.py
index 7178d7df84a12856bd3a6888d8bc0bad0f46e422..8174e7f50ccda4d0d1ed2496e8cce46760054fd9 100644
--- a/appengine/findit/waterfall/test/trigger_swarming_task_pipeline_test.py
+++ b/appengine/findit/waterfall/test/trigger_base_swarming_task_result_pipeline_test.py
@@ -5,16 +5,20 @@
 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'
@@ -23,10 +27,6 @@ class TriggerSwarmingTaskPipelineTest(wf_testcase.WaterfallTestCase):
     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
@@ -45,21 +45,21 @@ class TriggerSwarmingTaskPipelineTest(wf_testcase.WaterfallTestCase):
     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)
@@ -106,7 +106,7 @@ class TriggerSwarmingTaskPipelineTest(wf_testcase.WaterfallTestCase):
 
     def MockedGetSwarmingTaskName(*_):
       return 'new_task_name'
-    self.mock(trigger_swarming_task_pipeline, '_GetSwarmingTaskName',
+    self.mock(TriggerBaseSwarmingTaskPipeline, '_GetSwarmingTaskName',
               MockedGetSwarmingTaskName)
 
     master_name = 'm'