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