[Findit] Flake Checker: Get swarming task triggered by flake try jobs.

appengine/findit/waterfall/flake/test/process_flake_try_job_result_pipeline_test.py

Index: appengine/findit/waterfall/flake/test/process_flake_try_job_result_pipeline_test.py
diff --git a/appengine/findit/waterfall/flake/test/process_flake_try_job_result_pipeline_test.py b/appengine/findit/waterfall/flake/test/process_flake_try_job_result_pipeline_test.py
index 56770e3d909ac60efdd8a983ea8a2187129f61e7..1932fd168d36cc6c334fa5ccd15f8ae2a35b5d3c 100644
--- a/appengine/findit/waterfall/flake/test/process_flake_try_job_result_pipeline_test.py
+++ b/appengine/findit/waterfall/flake/test/process_flake_try_job_result_pipeline_test.py
@@ -2,10 +2,13 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-from gae_libs.testcase import TestCase
+import mock
 
+from gae_libs.testcase import TestCase
 from model.flake.flake_try_job import FlakeTryJob
 from model.flake.master_flake_analysis import MasterFlakeAnalysis
+from waterfall import swarming_util
+from waterfall.flake import process_flake_try_job_result_pipeline
 from waterfall.flake.process_flake_try_job_result_pipeline import (
     ProcessFlakeTryJobResultPipeline)
 
@@ -22,7 +25,6 @@ class ProcessFlakeTryJobResultPipelineTest(TestCase):
     commit_position = 4
     try_job_id = 'try_job_id'
     url = 'url'
-
     try_job_result = {
         'report': {
             'result': {
@@ -42,11 +44,9 @@ class ProcessFlakeTryJobResultPipelineTest(TestCase):
             }
         }
     }
-
     analysis = MasterFlakeAnalysis.Create(
         master_name, builder_name, build_number, step_name, test_name)
     analysis.Save()
-
     try_job = FlakeTryJob.Create(
         master_name, builder_name, step_name, test_name, revision)
     try_job.flake_results = [{
@@ -56,11 +56,9 @@ class ProcessFlakeTryJobResultPipelineTest(TestCase):
     }]
     try_job.try_job_ids = [try_job_id]
     try_job.put()
-
     ProcessFlakeTryJobResultPipeline().run(
         revision, commit_position, try_job_result, try_job.key.urlsafe(),
         analysis.key.urlsafe())
-
     resulting_data_point = analysis.data_points[-1]
     self.assertEqual(0.2, resulting_data_point.pass_rate)
     self.assertEqual(commit_position, resulting_data_point.commit_position)
@@ -76,7 +74,6 @@ class ProcessFlakeTryJobResultPipelineTest(TestCase):
     commit_position = 4
     try_job_id = 'try_job_id'
     url = 'url'
-
     try_job_result = {
         'report': {
             'result': {
@@ -89,11 +86,9 @@ class ProcessFlakeTryJobResultPipelineTest(TestCase):
             }
         }
     }
-
     analysis = MasterFlakeAnalysis.Create(
         master_name, builder_name, build_number, step_name, test_name)
     analysis.Save()
-
     try_job = FlakeTryJob.Create(
         master_name, builder_name, step_name, test_name, revision)
     try_job.flake_results = [{
@@ -103,12 +98,172 @@ class ProcessFlakeTryJobResultPipelineTest(TestCase):
     }]
     try_job.try_job_ids = [try_job_id]
     try_job.put()
-
     ProcessFlakeTryJobResultPipeline().run(
         revision, commit_position, try_job_result, try_job.key.urlsafe(),
         analysis.key.urlsafe())
-
     resulting_data_point = analysis.data_points[-1]
     self.assertEqual(-1, resulting_data_point.pass_rate)
     self.assertEqual(commit_position, resulting_data_point.commit_position)
     self.assertEqual(url, resulting_data_point.try_job_url)
+
+  @mock.patch.object(swarming_util, 'GetIsolatedOutputForTask')
+  def testGetSwarmingTaskIdForTryJob(self, mock_fn):
+    output_json_1 = {
+        'per_iteration_data': [{}, {}]
+    }
+    output_json_2 = {
+        'per_iteration_data': [
+            {
+                'Test.One': 'log for Test.One'
+            }
+        ]
+    }
+    mock_fn.side_effect = [output_json_1, output_json_2]
+
+    revision = 'r0'
+    step_name = 'gl_tests'
+    test_name = 'Test.One'
+    report = {
+        'result': {
+            'r0': {
+                'gl_tests': {
+                    'status': 'passed',
+                    'valid': True,
+                    'pass_fail_counts': {
+                        'Test.One': {
+                            'pass_count': 100,
+                            'fail_count': 0
+                        }
+                    },
+                    'step_metadata': {
+                        'swarm_task_ids': ['task1', 'task2']
+
+                    }
+                }
+            }
+        }
+    }
+
+    task_id = process_flake_try_job_result_pipeline._GetSwarmingTaskIdForTryJob(
+        report, revision, step_name, test_name)
+    self.assertEquals( 'task2', task_id)
+
+  def testGetSwarmingTaskIdForTryJobNoReport(self):
+    self.assertIsNone(
+        process_flake_try_job_result_pipeline._GetSwarmingTaskIdForTryJob(
+            None, None, None, None))
+
+  @mock.patch.object(swarming_util, 'GetIsolatedOutputForTask')
+  def testGetSwarmingTaskIdForTryJobNotFoundTaskWithResult(self, mock_fn):
+    output_json = {
+        'per_iteration_data': [{}, {}]
+    }
+    mock_fn.return_result = output_json
+
+    revision = 'r0'
+    step_name = 'gl_tests'
+    test_name = 'Test.One'
+    report = {
+        'result': {
+            'r0': {
+                'gl_tests': {
+                    'status': 'passed',
+                    'valid': True,
+                    'pass_fail_counts': {
+                        'Test.One': {
+                            'pass_count': 100,
+                            'fail_count': 0
+                        }
+                    },
+                    'step_metadata': {
+                        'swarm_task_ids': ['task1', 'task2']
+                    }
+                }
+            }
+        }
+    }
+
+    self.assertIsNone(
+        process_flake_try_job_result_pipeline._GetSwarmingTaskIdForTryJob(
+            report, revision, step_name, test_name))
+
+  @mock.patch.object(swarming_util, 'GetIsolatedOutputForTask',
+                     return_value=None)
+  def testGetSwarmingTaskIdForTryJobNoOutputJson(self, _):
+    revision = 'r0'
+    step_name = 'gl_tests'
+    test_name = 'Test.One'
+    report = {
+      'result': {
+        'r0': {
+          'gl_tests': {
+            'status': 'passed',
+            'valid': True,
+            'pass_fail_counts': {
+              'Test.One': {
+                'pass_count': 100,
+                'fail_count': 0
+              }
+            },
+            'step_metadata': {
+              'swarm_task_ids': ['task1', 'task2']
+            }
+          }
+        }
+      }
+    }
+
+    self.assertIsNone(
+      process_flake_try_job_result_pipeline._GetSwarmingTaskIdForTryJob(
+        report, revision, step_name, test_name))
+
+  def testGetSwarmingTaskIdForTryJobOnlyOneTask(self):
+    revision = 'r0'
+    step_name = 'gl_tests'
+    test_name = 'Test.One'
+    report = {
+      'result': {
+        'r0': {
+          'gl_tests': {
+            'status': 'passed',
+            'valid': True,
+            'pass_fail_counts': {
+              'Test.One': {
+                'pass_count': 100,
+                'fail_count': 0
+              }
+            },
+            'step_metadata': {
+              'swarm_task_ids': ['task1']
+            }
+          }
+        }
+      }
+    }
+
+    self.assertEquals(
+      process_flake_try_job_result_pipeline._GetSwarmingTaskIdForTryJob(
+        report, revision, step_name, test_name), 'task1')
+
+  def testGetSwarmingTaskIdForTryJobTestNotExist(self):
+    revision = 'r0'
+    step_name = 'gl_tests'
+    test_name = 'Test.One'
+    report = {
+      'result': {
+        'r0': {
+          'gl_tests': {
+            'status': 'passed',
+            'valid': True,
+            'pass_fail_counts': {},
+            'step_metadata': {
+              'swarm_task_ids': ['task1']
+            }
+          }
+        }
+      }
+    }
+
+    self.assertEquals(
+      process_flake_try_job_result_pipeline._GetSwarmingTaskIdForTryJob(
+        report, revision, step_name, test_name), 'task1')