[Findit] Refactoring trigger swarming task pipelines

appengine/findit/waterfall/test/trigger_swarming_task_pipeline_test.py

Index: appengine/findit/waterfall/test/trigger_swarming_task_pipeline_test.py
diff --git a/appengine/findit/waterfall/test/trigger_swarming_task_pipeline_test.py b/appengine/findit/waterfall/test/trigger_swarming_task_pipeline_test.py
index 7178d7df84a12856bd3a6888d8bc0bad0f46e422..44bac4e4c7e76cd025fe3821b0e47f917166df39 100644
--- a/appengine/findit/waterfall/test/trigger_swarming_task_pipeline_test.py
+++ b/appengine/findit/waterfall/test/trigger_swarming_task_pipeline_test.py
@@ -2,160 +2,63 @@
 # 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.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_swarming_task_pipeline import TriggerSwarmingTaskPipeline
 
 
 class TriggerSwarmingTaskPipelineTest(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']
-
-    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()
+    tests = []
 
-    pipeline = TriggerSwarmingTaskPipeline()
-    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),
+        TriggerSwarmingTaskPipeline()._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 = WfSwarmingTask.Create(
-        master_name, builder_name, build_number, step_name)
-    swarming_task.status = analysis_status.PENDING
-    swarming_task.put()
-
-    def MockedSleep(*_):
-      swarming_task = WfSwarmingTask.Get(
-          master_name, builder_name, build_number, step_name)
-      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()
-    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)
+    WfSwarmingTask.Create(
+        master_name, builder_name, build_number, step_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)
-
-    new_request_json = {}
-    def MockedTriggerSwarmingTask(new_request, *_):
-      new_request_json.update(new_request.Serialize())
-      return 'new_task_id'
-    self.mock(swarming_util, 'TriggerSwarmingTask', MockedTriggerSwarmingTask)
+    task = TriggerSwarmingTaskPipeline()._GetSwarmingTask(
+        master_name, builder_name, build_number, step_name)
 
-    def MockedGetSwarmingTaskName(*_):
-      return 'new_task_name'
-    self.mock(trigger_swarming_task_pipeline, '_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)
 
+  def testCreateSwarmingTask(self):
     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,
-        '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=10',
-                '--test-launcher-retry-limit=0',
-                '--gtest_filter=a.b:a.c'],
-            '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 = TriggerSwarmingTaskPipeline()
-    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)
+    build_number = 123
+    step_name = 's'
 
-    swarming_task = WfSwarmingTask.Get(
+    task = TriggerSwarmingTaskPipeline()._CreateSwarmingTask(
         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(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)
+
+  def testGetIterationsToRerun(self):
+    expected_iterations = 50
+    self.UpdateUnitTestConfigSettings(
+        config_property='swarming_settings',
+        override_data={'iterations_to_rerun': expected_iterations})
     self.assertEqual(
-        waterfall_config.GetSwarmingSettings()['iterations_to_rerun'],
-        swarming_task.parameters['iterations_to_rerun'])
+        expected_iterations,
+        TriggerSwarmingTaskPipeline()._GetIterationsToRerun())
+
+
+