Waterfall components of regression range finder.

appengine/findit/waterfall/test/trigger_flake_swarming_task_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
+import json
 
 from model import analysis_status
-from model.wf_swarming_task import WfSwarmingTask
+from model.flake.flake_swarming_task import FlakeSwarmingTask
 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
+from waterfall.trigger_flake_swarming_task_pipeline import (
+    TriggerFlakeSwarmingTaskPipeline)
 
 
-class TriggerSwarmingTaskPipelineTest(wf_testcase.WaterfallTestCase):
+class TriggerFlakeSwarmingTaskPipelineTest(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 = 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()
 
-    pipeline = TriggerSwarmingTaskPipeline()
+    pipeline = TriggerFlakeSwarmingTaskPipeline()
     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(TriggerFlakeSwarmingTaskPipeline, '_GetSwarmingTaskName',
               MockedGetSwarmingTaskName)
 
     master_name = 'm'
     builder_name = 'b'
     build_number = 234
     step_name = 'a_tests on platform'
-    tests = ['a.b', 'a.c']
+    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=10',
                 '--test-launcher-retry-limit=0',
-                '--gtest_filter=a.b:a.c'],
+                '--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 = TriggerSwarmingTaskPipeline()
+    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 = 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.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'],
+        pipeline._GetIterationsToRerun(),
         swarming_task.parameters['iterations_to_rerun'])