[Findit] Adding suspected CL to master flake analysis in preparation for try jobs

appengine/findit/model/flake/test/master_flake_analysis_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.
 
 from datetime import datetime
 
-import unittest
+from gae_libs.testcase import TestCase
 
 from model import analysis_status
 from model import result_status
 from model import triage_status
+from model.flake.flake_culprit import FlakeCulprit
 from model.flake.master_flake_analysis import DataPoint
 from model.flake.master_flake_analysis import MasterFlakeAnalysis
 
 
-class MasterFlakeAnalysisTest(unittest.TestCase):
+class MasterFlakeAnalysisTest(TestCase):
 
   def testMasterFlakeAnalysisStatusIsCompleted(self):
     for status in (analysis_status.COMPLETED, analysis_status.ERROR):
       analysis = MasterFlakeAnalysis.Create('m', 'b', 123, 's', 't')
       analysis.status = status
       self.assertTrue(analysis.completed)
 
   def testMasterFlakeAnalysisStatusIsNotCompleted(self):
     for status in (analysis_status.PENDING, analysis_status.RUNNING):
       analysis = MasterFlakeAnalysis.Create('m', 'b', 123, 's', 't')
@@ skipping 75 matching lines @@
 
   def testReset(self):
     analysis = MasterFlakeAnalysis.Create('m', 'b', 123, 's', 't')
     analysis.swarming_rerun_results = [{}]
     analysis.status = analysis_status.RUNNING
     analysis.correct_regression_range = True
     analysis.correct_culprit = False
     analysis.correct_culprit = None
     analysis.data_points = [DataPoint()]
     analysis.suspected_flake_build_number = 123
+    analysis.culprit = FlakeCulprit.Create('r', 'a1b2c3d4', 12345, 'url')
+    analysis.try_job_status = analysis_status.COMPLETED
     analysis.Reset()
 
     self.assertEqual([], analysis.swarming_rerun_results)
     self.assertEqual(analysis_status.PENDING, analysis.status)
     self.assertIsNone(analysis.correct_regression_range)
     self.assertIsNone(analysis.correct_culprit)
     self.assertIsNone(analysis.suspected_flake_build_number)
     self.assertEqual([], analysis.data_points)
+    self.assertIsNone(analysis.culprit)
+    self.assertIsNone(analysis.try_job_status)
 
   def testGetErrorMessage(self):
     cases = [
         (None, None),
         ('error', {'message': 'error', 'code': 'code'}),
     ]
     for expected_error_message, error in cases:
       analysis = MasterFlakeAnalysis.Create('m', 'b', 123, 's', 't')
       analysis.error = error
       self.assertEqual(expected_error_message, analysis.error_message)
@@ skipping 17 matching lines @@
 
     key = MasterFlakeAnalysis.Create(
         master_name, builder_name, build_number, step_name, test_name).key
 
     self.assertEqual(
         (None, None),
         MasterFlakeAnalysis.GetBuildConfigurationFromKey(None))
     self.assertEqual(
         (master_name, builder_name),
         MasterFlakeAnalysis.GetBuildConfigurationFromKey(key))
+
+  def testGetDataPointOfSuspectedBuildNoSuspectedFlakeBuildNumber(self):
+    analysis = MasterFlakeAnalysis.Create('m', 'b', 123, 's', 't')
+    self.assertIsNone(analysis.GetDataPointOfSuspectedBuild())
+
+  def testGetDataPointOfSuspectedBuild(self):
+    expected_build_number = 123
+    data_point = DataPoint()
+    data_point.build_number = expected_build_number
+
+    analysis = MasterFlakeAnalysis.Create('m', 'b', 125, 's', 't')
+    analysis.suspected_flake_build_number = expected_build_number
+    analysis.data_points.append(data_point)
+
+    suspected_data_point = analysis.GetDataPointOfSuspectedBuild()
+    self.assertEqual(expected_build_number, suspected_data_point.build_number)
+
+  def testGetDataPointOfSuspectedBuildNoDatapoint(self):
+    # This scenario should not happen.
+    expected_build_number = 123
+    unexpected_build_number = 124
+    data_point = DataPoint()
+    data_point.build_number = expected_build_number
+
+    analysis = MasterFlakeAnalysis.Create('m', 'b', 125, 's', 't')
+    analysis.suspected_flake_build_number = unexpected_build_number
+    analysis.data_points.append(data_point)
+
+    self.assertIsNone(analysis.GetDataPointOfSuspectedBuild())
+
+