[Findit] Refactor Findit pipeline.

appengine/findit/waterfall/test/schedule_try_job_pipeline_test.py

Index: appengine/findit/waterfall/test/schedule_try_job_pipeline_test.py
diff --git a/appengine/findit/waterfall/test/schedule_try_job_pipeline_test.py b/appengine/findit/waterfall/test/schedule_try_job_pipeline_test.py
index 8f124aa9e46bc2b6ffa513fdcc09630c20d4200e..d6ea7b51da06b9b487d61803110577b96bbf974e 100644
--- a/appengine/findit/waterfall/test/schedule_try_job_pipeline_test.py
+++ b/appengine/findit/waterfall/test/schedule_try_job_pipeline_test.py
@@ -2,30 +2,20 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import mock
+
 from common.waterfall import buildbucket_client
+from common.waterfall import failure_type
 from model.wf_try_job import WfTryJob
 from model.wf_try_job_data import WfTryJobData
+from waterfall import schedule_try_job_pipeline
 from waterfall.schedule_try_job_pipeline import ScheduleTryJobPipeline
 from waterfall.test import wf_testcase
-from waterfall.try_job_type import TryJobType
 
 
 class ScheduleTryjobPipelineTest(wf_testcase.WaterfallTestCase):
 
-  def _Mock_TriggerTryJobs(self, responses):
-    def MockedTriggerTryJobs(*_):
-      results = []
-      for response in responses:
-        if response.get('error'):  # pragma: no cover
-          results.append((
-              buildbucket_client.BuildbucketError(response['error']), None))
-        else:
-          results.append((
-              None, buildbucket_client.BuildbucketBuild(response['build'])))
-      return results
-    self.mock(buildbucket_client, 'TriggerTryJobs', MockedTriggerTryJobs)
-
-  def testGetBuildPropertiesWithCompileTargets(self):
+  def testGetBuildPropertiesWithSuspectedRevision(self):
     master_name = 'm'
     builder_name = 'b'
     build_number = 1
@@ -35,18 +25,19 @@ class ScheduleTryjobPipelineTest(wf_testcase.WaterfallTestCase):
         'good_revision': 1,
         'bad_revision': 2,
         'target_mastername': master_name,
-        'target_buildername': 'b',
         'referenced_build_url': ('https://build.chromium.org/p/%s/builders'
                                  '/%s/builds/%s') % (
-                                     master_name, builder_name, build_number)
+                                     master_name, builder_name, build_number),
+        'suspected_revisions': ['rev']
     }
     try_job_pipeline = ScheduleTryJobPipeline()
     properties = try_job_pipeline._GetBuildProperties(
-        master_name, builder_name, build_number, 1, 2, TryJobType.COMPILE, None)
+        master_name, builder_name, build_number, 1, 2, failure_type.COMPILE,
+        ['rev'])
 
     self.assertEqual(properties, expected_properties)
 
-  def testGetBuildPropertiesForTestFailure(self):
+  def testGetBuildPropertiesNoSuspectedRevision(self):
     master_name = 'm'
     builder_name = 'b'
     build_number = 1
@@ -56,96 +47,54 @@ class ScheduleTryjobPipelineTest(wf_testcase.WaterfallTestCase):
         'good_revision': 1,
         'bad_revision': 2,
         'target_mastername': master_name,
-        'target_testername': 'b',
         'referenced_build_url': ('https://build.chromium.org/p/%s/builders'
                                  '/%s/builds/%s') % (
                                      master_name, builder_name, build_number)
     }
     try_job_pipeline = ScheduleTryJobPipeline()
     properties = try_job_pipeline._GetBuildProperties(
-        master_name, builder_name, build_number, 1, 2, TryJobType.TEST, None)
+        master_name, builder_name, build_number, 1, 2, failure_type.TEST, None)
 
     self.assertEqual(properties, expected_properties)
 
-  def testSuccessfullyScheduleNewTryJobForCompileWithSuspectedRevisions(self):
+  @mock.patch.object(schedule_try_job_pipeline, 'buildbucket_client')
+  def testTriggerTryJob(self, mock_other):
     master_name = 'm'
     builder_name = 'b'
-    build_number = 223
-    good_revision = 'rev1'
-    bad_revision = 'rev2'
-    build_id = '1'
-    url = 'url'
-
-    responses = [
-        {
-            'build': {
-                'id': build_id,
-                'url': url,
-                'status': 'SCHEDULED',
-            }
+    response = {
+        'build': {
+            'id': '1',
+            'url': 'url',
+            'status': 'SCHEDULED',
         }
-    ]
+    }
+    results = [(None, buildbucket_client.BuildbucketBuild(response['build']))]
+    mock_other.TriggerTryJobs.return_value = results
 
-    self._Mock_TriggerTryJobs(responses)
+    try_job_pipeline = ScheduleTryJobPipeline()
+    build_id = try_job_pipeline._TriggerTryJob(
+        master_name, builder_name, {}, [])
 
-    WfTryJob.Create(master_name, builder_name, build_number).put()
+    self.assertEqual(build_id, '1')
 
-    try_job_pipeline = ScheduleTryJobPipeline()
-    try_job_id = try_job_pipeline.run(
-        master_name, builder_name, build_number, good_revision, bad_revision,
-        TryJobType.COMPILE, None, None, ['r5'])
-
-    try_job = WfTryJob.Get(master_name, builder_name, build_number)
-    try_job_data = WfTryJobData.Get(build_id)
-
-    expected_try_job_id = '1'
-    self.assertEqual(expected_try_job_id, try_job_id)
-    self.assertEqual(
-        expected_try_job_id, try_job.compile_results[-1]['try_job_id'])
-    self.assertTrue(expected_try_job_id in try_job.try_job_ids)
-    self.assertIsNotNone(try_job_data)
-    self.assertEqual(try_job_data.master_name, master_name)
-    self.assertEqual(try_job_data.builder_name, builder_name)
-    self.assertEqual(try_job_data.build_number, build_number)
-    self.assertEqual(try_job_data.try_job_type, TryJobType.COMPILE)
-    self.assertFalse(try_job_data.has_compile_targets)
-    self.assertTrue(try_job_data.has_heuristic_results)
-
-  def testSuccessfullyScheduleNewTryJobForTest(self):
+
+  def testCreateTryJobData(self):
+    build_id = 1
     master_name = 'm'
     builder_name = 'b'
-    build_number = 223
-    good_revision = 'rev1'
-    bad_revision = 'rev2'
-    targeted_tests = {'a': ['test1', 'test2']}
-    build_id = '1'
-
-    responses = [
-        {
-            'build': {
-                'id': build_id,
-                'url': 'url',
-                'status': 'SCHEDULED',
-            }
-        }
-    ]
-    self._Mock_TriggerTryJobs(responses)
+    build_number = 1
+    try_job_type = failure_type.GetDescriptionForFailureType(failure_type.TEST)
 
-    WfTryJob.Create(master_name, builder_name, build_number).put()
+    ScheduleTryJobPipeline()._CreateTryJobData(
+        build_id, master_name, builder_name, build_number, try_job_type,
+        False, False)
 
-    try_job_pipeline = ScheduleTryJobPipeline()
-    try_job_id = try_job_pipeline.run(
-        master_name, builder_name, build_number, good_revision, bad_revision,
-        TryJobType.TEST, None, targeted_tests, None)
-
-    try_job = WfTryJob.Get(master_name, builder_name, build_number)
-    try_job_data = WfTryJobData.Get(try_job_id)
-    self.assertEqual(try_job_id, build_id)
-    self.assertEqual(try_job.test_results[-1]['try_job_id'], build_id)
-    self.assertIsNotNone(try_job_data)
-    self.assertEqual(try_job_data.master_name, master_name)
-    self.assertEqual(try_job_data.builder_name, builder_name)
-    self.assertEqual(try_job_data.build_number, build_number)
-    self.assertEqual(try_job_data.try_job_type, TryJobType.TEST)
-    self.assertFalse(try_job_data.has_compile_targets)
-    self.assertFalse(try_job_data.has_heuristic_results)
+    data = WfTryJobData.Get(build_id)
+
+    self.assertIsNotNone(data)
+    self.assertEqual(data.master_name, master_name)
+    self.assertEqual(data.builder_name, builder_name)
+    self.assertEqual(data.build_number, build_number)
+    self.assertEqual(data.try_job_type, try_job_type)
+    self.assertFalse(data.has_compile_targets)
+    self.assertFalse(data.has_heuristic_results)