[Findit] Refactoring WfTryJobData into BaseTryJobData, WfTryJobData, and FlakeTryJobData

appengine/findit/waterfall/test/monitor_try_job_pipeline_test.py

 # Copyright 2015 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.
 
 from datetime import datetime
 import json
 import mock
 import time
 
 from common.waterfall import buildbucket_client
+from common.waterfall import failure_type
 from common.waterfall import try_job_error
-from common.waterfall import failure_type
 from model import analysis_status
 from model.flake.flake_try_job import FlakeTryJob
+from model.flake.flake_try_job_data import FlakeTryJobData
 from model.wf_try_job import WfTryJob
 from model.wf_try_job_data import WfTryJobData
 from waterfall import monitor_try_job_pipeline
 from waterfall import waterfall_config
 from waterfall.monitor_try_job_pipeline import MonitorTryJobPipeline
 from waterfall.test import wf_testcase
 
 
 class MonitorTryJobPipelineTest(wf_testcase.WaterfallTestCase):
 
@@ skipping 27 matching lines @@
         {'a': 1},
         {'a': 1, 'b': 2}))
 
   def testUpdateTryJobMetadataForBuildError(self):
     error_data = {
         'reason': 'BUILD_NOT_FOUND',
         'message': 'message'
     }
     error = buildbucket_client.BuildbucketError(error_data)
     try_job_data = WfTryJobData.Create('1')
+    try_job_data.try_job_key = WfTryJob.Create('m', 'b', 123).key
 
     monitor_try_job_pipeline._UpdateTryJobMetadata(
-        try_job_data, None, None, error, False)
+        try_job_data, failure_type.COMPILE, None, None, error, False)
     self.assertEqual(try_job_data.error, error_data)
 
   def testUpdateTryJobMetadata(self):
     try_job_id = '1'
     url = 'url'
     build_data = {
         'id': try_job_id,
         'url': url,
         'status': 'COMPLETED',
         'completed_ts': '1454367574000000',
@@ skipping 13 matching lines @@
         })
     }
     build = buildbucket_client.BuildbucketBuild(build_data)
     expected_error_dict = {
         'message': 'Try job monitoring was abandoned.',
         'reason': ('Timeout after %s hours' %
                    waterfall_config.GetTryJobSettings().get(
                        'job_timeout_hours'))
     }
     try_job_data = WfTryJobData.Create(try_job_id)
+    try_job_data.try_job_key = WfTryJob.Create('m', 'b', 123).key
 
     monitor_try_job_pipeline._UpdateTryJobMetadata(
-        try_job_data, None, build, None, False)
+        try_job_data, failure_type.COMPILE, None, build, None, False)
     try_job_data = WfTryJobData.Get(try_job_id)
     self.assertIsNone(try_job_data.error)
     self.assertEqual(try_job_data.regression_range_size, 2)
     self.assertEqual(try_job_data.number_of_commits_analyzed, 2)
     self.assertEqual(try_job_data.end_time, datetime(2016, 2, 1, 22, 59, 34))
     self.assertEqual(try_job_data.request_time,
                      datetime(2016, 2, 1, 22, 59, 30))
     self.assertEqual(try_job_data.try_job_url, url)
 
     monitor_try_job_pipeline._UpdateTryJobMetadata(
-        try_job_data, None, build, None, True)
+        try_job_data, failure_type.COMPILE, None, build, None, True)
     self.assertEqual(try_job_data.error, expected_error_dict)
     self.assertEqual(try_job_data.error_code, try_job_error.TIMEOUT)
 
   @mock.patch.object(monitor_try_job_pipeline, 'buildbucket_client')
   def testGetTryJobsForCompileSuccess(self, mock_module):
     master_name = 'm'
     builder_name = 'b'
     build_number = 1
     try_job_id = '1'
     regression_range_size = 2
 
     try_job = WfTryJob.Create(master_name, builder_name, build_number)
     try_job_data = WfTryJobData.Create(try_job_id)
+    try_job_data.try_job_key = try_job.key
     try_job_data.put()
     try_job.compile_results = [
         {
             'report': None,
             'url': 'url',
             'try_job_id': '1',
         }
     ]
     try_job.status = analysis_status.RUNNING
     try_job.put()
@@ skipping 58 matching lines @@
         {
             'report': None,
             'url': 'url',
             'try_job_id': try_job_id,
         }
     ]
     try_job.status = analysis_status.RUNNING
     try_job.put()
 
     try_job_data = WfTryJobData.Create(try_job_id)
+    try_job_data.try_job_key = try_job.key
     try_job_data.put()
 
     data = [
         {
             'build': {
                 'id': '3',
                 'url': 'url',
                 'status': 'STARTED'
             }
         },
@@ skipping 78 matching lines @@
         },
         'url': 'url',
         'try_job_id': '3',
     }
     self.assertEqual(expected_test_result, test_result)
 
     try_job = WfTryJob.Get(master_name, builder_name, build_number)
     self.assertEqual(expected_test_result, try_job.test_results[-1])
     self.assertEqual(analysis_status.RUNNING, try_job.status)
 
+  @mock.patch.object(monitor_try_job_pipeline, 'buildbucket_client')
+  def testGetTryJobsForFlakeSuccess(self, mock_module):
+    master_name = 'm'
+    builder_name = 'b'
+    step_name = 's'
+    test_name = 't'
+    git_hash = 'a1b2c3d4'
+    try_job_id = '1'
+
+    try_job = FlakeTryJob.Create(
+        master_name, builder_name, step_name, test_name, git_hash)
+    try_job.flake_results = [
+        {
+            'report': None,
+            'url': 'url',
+            'try_job_id': '1',
+        }
+    ]
+    try_job.status = analysis_status.RUNNING
+    try_job.put()
+
+    try_job_data = FlakeTryJobData.Create(try_job_id)
+    try_job_data.try_job_key = try_job.key
+    try_job_data.put()
+
+    build_response = {
+        'id': '1',
+        'url': 'url',
+        'status': 'COMPLETED',
+        'result_details_json': json.dumps({
+            'properties': {
+                'report': {
+                    'result': {
+                        'r0': {
+                            'gl_tests': {
+                                'status': 'passed',
+                                'valid': True,
+                                'pass_fail_counts': {
+                                    'Test.One': {
+                                        'pass_count': 100,
+                                        'fail_count': 0
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        })
+    }
+
+    mock_module.GetTryJobs.return_value = [
+        (None, buildbucket_client.BuildbucketBuild(build_response))]
+
+    pipeline = MonitorTryJobPipeline()
+    flake_result = pipeline.run(
+        try_job.key.urlsafe(), failure_type.FLAKY_TEST, try_job_id)
+
+    expected_flake_result = {
+        'report': {
+            'result': {
+                'r0': {
+                    'gl_tests': {
+                        'status': 'passed',
+                        'valid': True,
+                        'pass_fail_counts': {
+                            'Test.One': {
+                                'pass_count': 100,
+                                'fail_count': 0
+                            }
+                        }
+                    }
+                }
+            }
+        },
+        'url': 'url',
+        'try_job_id': '1',
+    }
+
+    self.assertEqual(expected_flake_result, flake_result)
+
+    try_job = FlakeTryJob.Get(
+        master_name, builder_name, step_name, test_name, git_hash)
+    self.assertEqual(expected_flake_result, try_job.flake_results[-1])
+    self.assertEqual(analysis_status.RUNNING, try_job.status)
+
+    try_job_data = FlakeTryJobData.Get(try_job_id)
+    self.assertEqual(try_job_data.last_buildbucket_response, build_response)
+
   def testUpdateTryJobResultAnalyzing(self):
     master_name = 'm'
     builder_name = 'b'
     build_number = 1
     try_job_id = '3'
 
     try_job = WfTryJob.Create(master_name, builder_name, build_number)
     try_job.put()
 
     pipeline = MonitorTryJobPipeline()
@@ skipping 182 matching lines @@
         (expected_error_dict, try_job_error.UNKNOWN))
 
   def testReturnNoneIfNoTryJobId(self):
     master_name = 'm'
     builder_name = 'b'
     build_number = 1
     pipeline = MonitorTryJobPipeline()
     try_job = WfTryJob.Create(master_name, builder_name, build_number)
     test_result = pipeline.run(try_job.key.urlsafe(), failure_type.TEST, None)
     self.assertIsNone(test_result)