[Findit] Refactoring trigger swarming task pipelines

appengine/findit/waterfall/test/trigger_flake_swarming_task_pipeline_test.py

Index: appengine/findit/waterfall/test/trigger_flake_swarming_task_pipeline_test.py
diff --git a/appengine/findit/waterfall/test/trigger_flake_swarming_task_pipeline_test.py b/appengine/findit/waterfall/test/trigger_flake_swarming_task_pipeline_test.py
index aa897ca16d647d06e0f5fc7d70bb4878411ec025..e6de1eca9c6559dd233364c29083c21acc0e7053 100644
--- a/appengine/findit/waterfall/test/trigger_flake_swarming_task_pipeline_test.py
+++ b/appengine/findit/waterfall/test/trigger_flake_swarming_task_pipeline_test.py
@@ -2,12 +2,7 @@
 # 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 waterfall import swarming_util
-from waterfall.swarming_task_request import SwarmingTaskRequest
 from waterfall.test import wf_testcase
 from waterfall.trigger_flake_swarming_task_pipeline import (
     TriggerFlakeSwarmingTaskPipeline)
@@ -15,142 +10,57 @@ from waterfall.trigger_flake_swarming_task_pipeline import (
 
 class TriggerFlakeSwarmingTaskPipelineTest(wf_testcase.WaterfallTestCase):
 
-  def testNoNewSwarmingTaskIsNeeded(self):
+  def testGetArgs(self):
     master_name = 'm'
     builder_name = 'b'
-    build_number = 1
+    build_number = 123
     step_name = 's'
-    tests = ['a.b']
-
-    swarming_task = FlakeSwarmingTask.Create(
-        master_name, builder_name, build_number, step_name, tests[0])
-    swarming_task.status = analysis_status.RUNNING
-    swarming_task.task_id = 'task_id'
-    swarming_task.put()
+    tests = ['t']
 
-    pipeline = TriggerFlakeSwarmingTaskPipeline()
-    task_id = pipeline.run(
-        master_name, builder_name, build_number, step_name, tests)
-    self.assertEqual('task_id', task_id)
+    self.assertEqual(
+        (master_name, builder_name, build_number, step_name, tests[0]),
+        TriggerFlakeSwarmingTaskPipeline()._GetArgs(
+            master_name, builder_name, build_number, step_name, tests))
 
-  def testWaitingForTheTaskId(self):
+  def testGetSwarmingTask(self):
     master_name = 'm'
     builder_name = 'b'
-    build_number = 1
+    build_number = 123
     step_name = 's'
-    tests = ['a.b']
-
-    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 = 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 = TriggerFlakeSwarmingTaskPipeline()
-    task_id = pipeline.run(
-        master_name, builder_name, build_number, step_name, tests)
-    self.assertEqual('task_id', task_id)
+    test_name = 't'
 
-  def testTriggerANewSwarmingTask(self):
-    def MockedDownloadSwarmingTaskData(*_):
-      return [{'task_id': '1'}, {'task_id': '2'}]
-    self.mock(swarming_util, 'ListSwarmingTasksDataByTags',
-              MockedDownloadSwarmingTaskData)
+    FlakeSwarmingTask.Create(
+        master_name, builder_name, build_number, step_name, test_name).put()
 
-    def MockedGetSwarmingTaskRequest(ref_task_id, *_):
-      self.assertEqual('1', ref_task_id)
-      return SwarmingTaskRequest.Deserialize({
-          'expiration_secs': 3600,
-          'name': 'ref_task_request',
-          'parent_task_id': 'pti',
-          'priority': 25,
-          'properties': {
-              'command': 'cmd',
-              'dimensions': [{'key': 'k', 'value': 'v'}],
-              'env': [
-                  {'key': 'a', 'value': '1'},
-                  {'key': 'GTEST_SHARD_INDEX', 'value': '1'},
-                  {'key': 'GTEST_TOTAL_SHARDS', 'value': '5'},
-              ],
-              'execution_timeout_secs': 3600,
-              'extra_args': ['--flag=value', '--gtest_filter=d.f'],
-              'grace_period_secs': 30,
-              'idempotent': True,
-              'inputs_ref': {'a': 1},
-              'io_timeout_secs': 1200,
-          },
-          'tags': ['master:a', 'buildername:b', 'name:a_tests'],
-          'user': 'user',
-      })
-    self.mock(swarming_util, 'GetSwarmingTaskRequest',
-              MockedGetSwarmingTaskRequest)
+    task = TriggerFlakeSwarmingTaskPipeline()._GetSwarmingTask(
+        master_name, builder_name, build_number, step_name, test_name)
 
-    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(TriggerFlakeSwarmingTaskPipeline, '_GetSwarmingTaskName',
-              MockedGetSwarmingTaskName)
+    self.assertEqual(master_name, task.master_name)
+    self.assertEqual(builder_name, task.builder_name)
+    self.assertEqual(build_number, task.build_number)
+    self.assertEqual(step_name, task.step_name)
+    self.assertEqual(test_name, task.test_name)
 
+  def testCreateSwarmingTask(self):
     master_name = 'm'
     builder_name = 'b'
-    build_number = 234
-    step_name = 'a_tests on platform'
-    tests = ['a.b']
-
-    expected_new_request_json = {
-        'expiration_secs': 3600,
-        'name': 'new_task_name',
-        'parent_task_id': '',
-        'priority': 25,
-        'properties': {
-            'command': 'cmd',
-            'dimensions': [{'key': 'k', 'value': 'v'}],
-            'env': [
-                {'key': 'a', 'value': '1'},
-            ],
-            'execution_timeout_secs': 3600,
-            'extra_args': [
-                '--flag=value', '--gtest_repeat=100',
-                '--test-launcher-retry-limit=0',
-                '--gtest_filter=a.b'],
-            'grace_period_secs': 30,
-            'idempotent': False,
-            'inputs_ref': {'a': 1},
-            'io_timeout_secs': 1200,
-        },
-        'tags': ['purpose:deflake', 'ref_master:%s' % master_name,
-                 'ref_buildername:%s' % builder_name,
-                 'ref_buildnumber:%s' % build_number,
-                 'ref_stepname:%s' % step_name,
-                 'ref_task_id:1',
-                 'ref_name:a_tests'],
-        'user': '',
-    }
-
-    pipeline = TriggerFlakeSwarmingTaskPipeline()
-    new_task_id = pipeline.run(
-        master_name, builder_name, build_number, step_name, tests)
-    self.assertEqual('new_task_id', new_task_id)
-    self.assertEqual(expected_new_request_json, new_request_json)
-
-    swarming_task = FlakeSwarmingTask.Get(
-        master_name, builder_name, build_number, step_name, tests[0])
-    self.assertIsNotNone(swarming_task)
-    self.assertEqual('new_task_id', swarming_task.task_id)
-    self.assertEqual(tests, swarming_task.parameters['tests'])
+    build_number = 123
+    step_name = 's'
+    test_name = 't'
+
+    task = TriggerFlakeSwarmingTaskPipeline()._CreateSwarmingTask(
+        master_name, builder_name, build_number, step_name, test_name)
+    self.assertEqual(master_name, task.master_name)
+    self.assertEqual(builder_name, task.builder_name)
+    self.assertEqual(build_number, task.build_number)
+    self.assertEqual(step_name, task.step_name)
+    self.assertEqual(test_name, task.test_name)
+
+  def testGetIterationsToRerun(self):
+    expected_iterations = 50
+    self.UpdateUnitTestConfigSettings(
+        config_property='check_flake_settings',
+        override_data={'iterations_to_rerun': expected_iterations})
     self.assertEqual(
-        pipeline._GetIterationsToRerun(),
-        swarming_task.parameters['iterations_to_rerun'])
+        expected_iterations,
+        TriggerFlakeSwarmingTaskPipeline()._GetIterationsToRerun())