def testRunBenchmarkTask(self, quota_flag_value, capacity_flag_value,
retry_count, run_results, expected_run_count,
expected_quota_retry_calls,
expected_capacity_retry_calls):
FLAGS.zones = ['test_zone']
FLAGS.cloud = 'GCP'
FLAGS.retries = retry_count
FLAGS.smart_quota_retry = quota_flag_value
FLAGS.smart_capacity_retry = capacity_flag_value
test_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml()
# Generate some exceptions for each run.
self.enter_context(
mock.patch.object(pkb, 'DoProvisionPhase',
side_effect=run_results))
# Mock the retry options.
mock_quota_retry = self.enter_context(
mock.patch.object(pkb.ZoneRetryManager, '_AssignZoneToNewRegion'))
mock_capacity_retry = self.enter_context(
mock.patch.object(pkb.ZoneRetryManager, '_AssignNewZone'))
self._MockGcpUtils('GetZonesFromMachineType', return_value=set())
benchmark_specs, _ = pkb.RunBenchmarkTask(spec=test_spec)
self.assertEqual(len(benchmark_specs), expected_run_count)
# Retry preparation functions should have the right calls.
self.assertEqual(mock_quota_retry.call_count,
expected_quota_retry_calls)
self.assertEqual(mock_capacity_retry.call_count,
expected_capacity_retry_calls)
def testRunRetries(self):
test_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml()
# Generate some exceptions for each run.
self.enter_context(
mock.patch.object(pkb,
'DoProvisionPhase',
side_effect=errors.Benchmarks.QuotaFailure()))
benchmark_specs, _ = pkb.RunBenchmarkTask(spec=test_spec)
self.assertEqual(len(benchmark_specs), 4)