[Findit] Use new mock for unittests.

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 model import analysis_status
 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
 
 
-# A counter to enable different responses to requests in a loop.
-REQUEST_COUNTER = {
-    '1': 0,
-    '2': 0,
-    '3': 0
-}
-
-
 class MonitorTryJobPipelineTest(wf_testcase.WaterfallTestCase):
 
-  def _MockGetTryJobs(self, build_id):
-    def Mocked_GetTryJobs(*_):
-      data = {
-          '1': [
-              {
-                  'build': {
-                      'id': '1',
-                      'url': 'url',
-                      'status': 'COMPLETED',
-                      'result_details_json': json.dumps({
-                          'properties': {
-                              'report': {
-                                  'result': {
-                                      'rev1': 'passed',
-                                      'rev2': 'failed'
-                                  },
-                                  'metadata': {
-                                      'regression_range_size': 2
-                                  }
-                              }
-                          }
-                      })
-                  }
-              }
-          ],
-          '3': [
-              {
-                  'build': {
-                      'id': '3',
-                      'url': 'url',
-                      'status': 'STARTED'
-                  }
-              },
-              {
-                  'error': {
-                      'reason': 'BUILD_NOT_FOUND',
-                      'message': 'message',
-                  }
-              },
-              {
-                  'build': {
-                      'id': '3',
-                      'url': 'url',
-                      'status': 'STARTED'
-                  }
-              },
-              {
-                  'error': {
-                      'reason': 'BUILD_NOT_FOUND',
-                      'message': 'message',
-                  }
-              },
-              {
-                  'build': {
-                      'id': '3',
-                      'url': 'url',
-                      'status': 'COMPLETED',
-                      'result_details_json': json.dumps({
-                          'properties': {
-                              'report': {
-                                  'result': {
-                                      'rev1': {
-                                          'a_test': {
-                                              'status': 'passed',
-                                              'valid': True
-                                          }
-                                      },
-                                      'rev2': {
-                                          'a_test': {
-                                              'status': 'failed',
-                                              'valid': True,
-                                              'failures': ['test1', 'test2']
-                                          }
-                                      }
-                                  }
-                              }
-                          }
-                      })
-                  }
-              }
-          ]
-      }
-      try_job_results = []
-      build_error = data.get(build_id)[REQUEST_COUNTER[build_id]]
-      if build_error.get('error'):
-        try_job_results.append((
-            buildbucket_client.BuildbucketError(build_error['error']), None))
-      else:
-        try_job_results.append((
-            None, buildbucket_client.BuildbucketBuild(build_error['build'])))
-      REQUEST_COUNTER[build_id] += 1
-      return try_job_results
-
-    self.mock(buildbucket_client, 'GetTryJobs', Mocked_GetTryJobs)
-
   def setUp(self):
     super(MonitorTryJobPipelineTest, self).setUp()
     self.mock(time, 'sleep', lambda x: None)
 
   def testMicrosecondsToDatetime(self):
     self.assertEqual(
         datetime(2016, 2, 1, 22, 59, 34),
         monitor_try_job_pipeline._MicrosecondsToDatetime(1454367574000000))
     self.assertIsNone(monitor_try_job_pipeline._MicrosecondsToDatetime(None))
 
@@ skipping 50 matching lines @@
     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)
     self.assertEqual(try_job_data.error, expected_error_dict)
     self.assertEqual(try_job_data.error_code, try_job_error.TIMEOUT)
 
-  def testGetTryJobsForCompileSuccess(self):
+  @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.put()
     try_job.compile_results = [
         {
             'report': None,
             'url': 'url',
             'try_job_id': '1',
         }
     ]
     try_job.status = analysis_status.RUNNING
     try_job.put()
-    self._MockGetTryJobs(try_job_id)
+
+    build_response = {
+        'id': '1',
+        'url': 'url',
+        'status': 'COMPLETED',
+        'result_details_json': json.dumps({
+            'properties': {
+                'report': {
+                    'result': {
+                        'rev1': 'passed',
+                        'rev2': 'failed'
+                    },
+                    'metadata': {
+                        'regression_range_size': 2
+                    }
+                }
+            }
+        })
+    }
+    mock_module.GetTryJobs.return_value = [
+        (None, buildbucket_client.BuildbucketBuild(build_response))]
 
     pipeline = MonitorTryJobPipeline()
     compile_result = pipeline.run(
         master_name, builder_name, build_number, failure_type.COMPILE,
         try_job_id)
 
     expected_compile_result = {
         'report': {
             'result': {
                 'rev1': 'passed',
                 'rev2': 'failed'
             },
             'metadata': {
                 'regression_range_size': regression_range_size
             }
         },
         'url': 'url',
         'try_job_id': '1',
     }
 
     self.assertEqual(expected_compile_result, compile_result)
 
     try_job = WfTryJob.Get(master_name, builder_name, build_number)
     self.assertEqual(expected_compile_result, try_job.compile_results[-1])
     self.assertEqual(analysis_status.RUNNING, try_job.status)
 
     try_job_data = WfTryJobData.Get(try_job_id)
     self.assertEqual(try_job_data.regression_range_size, regression_range_size)
 
-  def testGetTryJobsForTestSuccess(self):
+  @mock.patch.object(monitor_try_job_pipeline, 'buildbucket_client')
+  def testGetTryJobsForTestSuccess(self, mock_module):
     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.test_results = [
         {
             '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.put()
 
-    self._MockGetTryJobs(try_job_id)
+    data = [
+        {
+            'build': {
+                'id': '3',
+                'url': 'url',
+                'status': 'STARTED'
+            }
+        },
+        {
+            'error': {
+                'reason': 'BUILD_NOT_FOUND',
+                'message': 'message',
+            }
+        },
+        {
+            'build': {
+                'id': '3',
+                'url': 'url',
+                'status': 'STARTED'
+            }
+        },
+        {
+            'error': {
+                'reason': 'BUILD_NOT_FOUND',
+                'message': 'message',
+            }
+        },
+        {
+            'build': {
+                'id': '3',
+                'url': 'url',
+                'status': 'COMPLETED',
+                'result_details_json': json.dumps({
+                    'properties': {
+                        'report': {
+                            'result': {
+                                'rev1': {
+                                    'a_test': {
+                                        'status': 'passed',
+                                        'valid': True
+                                    }
+                                },
+                                'rev2': {
+                                    'a_test': {
+                                        'status': 'failed',
+                                        'valid': True,
+                                        'failures': ['test1', 'test2']
+                                    }
+                                }
+                            }
+                        }
+                    }
+                })
+            }
+        }
+    ]
+
+    mock_module.GetTryJobs.side_effect = [
+        [(None, buildbucket_client.BuildbucketBuild(data[0]['build']))],
+        [(buildbucket_client.BuildbucketError(data[1]['error']), None)],
+        [(None, buildbucket_client.BuildbucketBuild(data[2]['build']))],
+        [(buildbucket_client.BuildbucketError(data[3]['error']), None)],
+        [(None, buildbucket_client.BuildbucketBuild(data[4]['build']))],
+    ]
 
     pipeline = MonitorTryJobPipeline()
     test_result = pipeline.run(
         master_name, builder_name, build_number, failure_type.TEST,
         try_job_id)
 
     expected_test_result = {
         'report': {
             'result': {
                 'rev1': {
@@ skipping 186 matching lines @@
     }
 
     expected_error_dict = {
         'message': 'No result report was found.',
         'reason': MonitorTryJobPipeline.UNKNOWN
     }
 
     self.assertEqual(
         monitor_try_job_pipeline._GetError(build_response, None, False),
         (expected_error_dict, try_job_error.UNKNOWN))