[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
new file mode 100644
index 0000000000000000000000000000000000000000..0798a04139432e45025c08eb2a22fbc4c2dabf5c
--- /dev/null
+++ b/appengine/findit/waterfall/flake/test/process_flake_try_job_result_pipeline_test.py
@@ -0,0 +1,126 @@
+# Copyright 2017 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 mock
+
+from common import constants
+from model.flake.flake_analysis_request import BuildStep
+from model.flake.flake_analysis_request import FlakeAnalysisRequest
+from waterfall.flake import flake_analysis_service
+from waterfall.flake import process_flake_try_job_result_pipeline
+from waterfall import swarming_util
+from waterfall.test import wf_testcase
+from waterfall.test_info import TestInfo
+
+
+class ProcessFlakeTryJobResultPipelineTest(wf_testcase.WaterfallTestCase):
+
+  @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'
+    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)
+    self.assertEquals( 'task2', task_id)
+
+  def testGetSwarmingTaskIdForTryJobNoReport(self):
+    self.assertIsNone(
+        process_flake_try_job_result_pipeline._GetSwarmingTaskIdForTryJob(
+            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'
+    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.assertIsNone(
+        process_flake_try_job_result_pipeline._GetSwarmingTaskIdForTryJob(
+            report, revision, step_name))
+
+
+  @mock.patch.object(swarming_util, 'GetIsolatedOutputForTask',
+                     return_value=None)
+  def testGetSwarmingTaskIdForTryJobNoOutputJson(self, _):
+    revision = 'r0'
+    step_name = 'gl_tests'
+    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.assertIsNone(
+      process_flake_try_job_result_pipeline._GetSwarmingTaskIdForTryJob(
+        report, revision, step_name))