def testShouldRetry(self, status, failed_substatus, retry_substatuses,
expected_retry):
test_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml()
test_spec.status = status
test_spec.failed_substatus = failed_substatus
with flagsaver.flagsaver(retry_substatuses=retry_substatuses):
self.assertEqual(pkb._ShouldRetry(test_spec), expected_retry)
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 testSimpleConfig(self):
spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
yaml_string=_SIMPLE_EDW_CONFIG, benchmark_name='edw_benchmark')
spec.ConstructEdwService()
self.assertEqual('snowflake_aws', spec.edw_service.SERVICE_TYPE)
self.assertIsInstance(spec.edw_service,
providers.aws.snowflake.Snowflake)
def testZonesFlag(self):
FLAGS.zones = ['us-east-1b', 'zone2']
FLAGS.extra_zones = []
spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
MULTI_CLOUD_CONFIG)
spec.ConstructVirtualMachines()
self.assertEqual(len(spec.vms), 2)
self.assertEqual(spec.vm_groups['group1'][0].zone, 'us-east-1b')
self.assertEqual(spec.vm_groups['group2'][0].zone, 'zone2')
def testValidConfigWithDiskSpec(self):
spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
VALID_CONFIG_WITH_DISK_SPEC)
spec.ConstructVirtualMachines()
vms = spec.vm_groups['default']
self.assertEqual(len(vms), 2)
for vm in vms:
self.assertEqual(len(vm.disk_specs), 3)
self.assertTrue(
all(disk_spec.disk_size == 75 for disk_spec in vm.disk_specs))
def testMultiCloud(self):
spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
MULTI_CLOUD_CONFIG)
spec.ConstructVirtualMachines()
self.assertEqual(len(spec.vms), 2)
self.assertIsInstance(spec.vm_groups['group1'][0],
aws_vm.AwsVirtualMachine)
self.assertIsInstance(spec.vm_groups['group2'][0],
gce_vm.GceVirtualMachine)
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)
def testSimpleConfig(self):
spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(SIMPLE_CONFIG)
spec.ConstructVirtualMachines()
self.assertEqual(len(spec.vms), 1)
vm = spec.vms[0]
self.assertEqual(vm.machine_type, 'n1-standard-4')
self.assertEqual(vm.zone, 'us-central1-c')
self.assertEqual(vm.project, 'my-project')
self.assertEqual(vm.disk_specs, [])
def setUp(self):
super().setUp()
test_spec = inspect.cleandoc("""
cloud_spanner_ycsb:
description: Sample spanner benchmark
spanner:
service_type: default
enable_freeze_restore: True
delete_on_freeze_error: True
create_on_restore_error: True
""")
self.test_bm_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
yaml_string=test_spec, benchmark_name='cloud_spanner_ycsb')
def testRunBenchmarkExceptionHasCorrectFailureStatus(
self, exception_class, expected_substatus):
self.enter_context(
mock.patch.object(pkb,
'DoProvisionPhase',
side_effect=[exception_class()]))
test_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml()
# Skip pickling the spec.
self.enter_context(mock.patch.object(test_spec, 'Pickle'))
with self.assertRaises(exception_class):
pkb.RunBenchmark(spec=test_spec, collector=mock.Mock())
self.assertEqual(test_spec.status, benchmark_status.FAILED)
self.assertEqual(test_spec.failed_substatus, expected_substatus)
def testRestoreInstanceCopiedFromPreviousSpec(self):
restore_spanner_spec = inspect.cleandoc("""
cloud_spanner_ycsb:
spanner:
name: restore_spanner
service_type: default
""")
# Set up the restore spec Spanner instance
self.test_bm_spec.restore_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
yaml_string=restore_spanner_spec,
benchmark_name='cloud_spanner_ycsb')
self.test_bm_spec.restore_spec.ConstructSpanner()
self.test_bm_spec.ConstructSpanner()
self.assertEqual(self.test_bm_spec.spanner.name, 'restore_spanner')
def testStaticVms(self):
spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
STATIC_VM_CONFIG)
spec.ConstructVirtualMachines()
self.assertEqual(len(spec.vms), 4)
vm0 = spec.vm_groups['group1'][0]
vm1, vm2, vm3 = spec.vm_groups['group2']
self.assertIsInstance(vm0, gce_vm.GceVirtualMachine)
self.assertIsInstance(vm1, static_vm.StaticVirtualMachine)
self.assertIsInstance(vm2, static_vm.Rhel7BasedStaticVirtualMachine)
self.assertIsInstance(vm3, gce_vm.GceVirtualMachine)
self.assertEqual(vm2.disk_specs[0].mount_point, '/scratch')
def testSmartCapacityRetry(self):
test_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml()
test_spec.failed_substatus = (
benchmark_status.FailedSubstatus.INSUFFICIENT_CAPACITY)
self._MockLoadProviderUtils(gcp_utils)
self._MockGcpUtils('GetRegionFromZone')
# Expect that the correct possible zones are passed to the function below.
self._MockGcpUtils('GetZonesInRegion',
return_value={'zone_1', 'zone_2'})
test_retry_manager = pkb.ZoneRetryManager()
test_retry_manager.HandleSmartRetries(test_spec)
# zone_1 is recorded.
self.assertEqual(test_retry_manager._zones_tried, {'zone_1'})
# zone_2 is the new zone picked.
self.assertEqual(FLAGS.zone, ['zone_2'])
def testCompletionStatusesContainFreezeRestoreStatus(self):
# Arrange
test_bm_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml()
test_bm_spec.status = benchmark_status.SUCCEEDED
mock_status_file = mock.Mock()
test_bm_spec.config.flags = {}
# Act
pkb._WriteCompletionStatusFile([test_bm_spec], mock_status_file)
# Assert
mock_status_file.write.assert_called_once_with(
json.dumps({
'name': 'cluster_boot',
'status': benchmark_status.SUCCEEDED,
'flags': {
'freeze': 'mock_freeze_path',
'restore': 'mock_restore_path'
}
}) + '\n')
def testSpecPickledEvenWithException(self):
# Make provision have an exception.
self.enter_context(
mock.patch.object(pkb, 'DoProvisionPhase',
side_effect=Exception()))
# Make benchmark_spec avoid pickling for test.
test_benchmark_spec = (
pkb_common_test_case.CreateBenchmarkSpecFromYaml(
pkb_common_test_case.SIMPLE_CONFIG, 'cluster_boot'))
self.enter_context(mock.patch.object(test_benchmark_spec, 'Pickle'))
mock_freeze = self.enter_context(
mock.patch.object(test_benchmark_spec, 'Freeze'))
collector = publisher.SampleCollector()
# Run the benchmark loop.
with self.assertRaises(Exception):
pkb.RunBenchmark(test_benchmark_spec, collector)
# PKB should still attempt to freeze benchmark spec.
mock_freeze.assert_called_once()
def testSmartQuotaRetry(self):
test_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml()
test_spec.failed_substatus = benchmark_status.FailedSubstatus.QUOTA
# Start with zone_1 in region_1.
self._MockLoadProviderUtils(gcp_utils)
self._MockGcpUtils('GetRegionFromZone', return_value='region_1')
self._MockGcpUtils('GetGeoFromRegion')
self._MockGcpUtils('GetRegionsInGeo',
return_value={'region_1', 'region_2'})
# Expect that region_1 is skipped when getting zones.
mock_get_zones = self._MockGcpUtils('GetZonesInRegion',
return_value={'zone_2'})
test_retry_manager = pkb.ZoneRetryManager()
test_retry_manager.HandleSmartRetries(test_spec)
# Function should not get zones from region_1, resulting in only 1 call.
mock_get_zones.assert_called_once()
# zone_1 is recorded.
self.assertEqual(test_retry_manager._zones_tried, {'zone_1'})
# zone_2 is the new zone picked.
self.assertEqual(FLAGS.zone, ['zone_2'])
def GetTestDynamoDBInstance(minimal=False):
spec = _MINIMAL_TEST_BENCHMARK_SPEC if minimal else _TEST_BENCHMARK_SPEC
test_benchmark_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
yaml_string=spec, benchmark_name='aws_dynamodb_ycsb')
test_benchmark_spec.ConstructNonRelationalDb()
return test_benchmark_spec.non_relational_db
def GetTestBigtableInstance(spec=_TEST_BENCHMARK_SPEC):
test_benchmark_spec = pkb_common_test_case.CreateBenchmarkSpecFromYaml(
yaml_string=spec, benchmark_name='cloud_bigtable_ycsb')
test_benchmark_spec.ConstructNonRelationalDb()
return test_benchmark_spec.non_relational_db
