1# Copyright (c) 2011-present, Facebook, Inc. All rights reserved. 2# This source code is licensed under both the GPLv2 (found in the 3# COPYING file in the root directory) and Apache 2.0 License 4# (found in the LICENSE.Apache file in the root directory). 5 6from advisor.db_bench_runner import DBBenchRunner 7from advisor.db_log_parser import NO_COL_FAMILY, DataSource 8from advisor.db_options_parser import DatabaseOptions 9import os 10import unittest 11 12 13class TestDBBenchRunnerMethods(unittest.TestCase): 14 def setUp(self): 15 self.pos_args = [ 16 './../../db_bench', 17 'overwrite', 18 'use_existing_db=true', 19 'duration=10' 20 ] 21 self.bench_runner = DBBenchRunner(self.pos_args) 22 this_path = os.path.abspath(os.path.dirname(__file__)) 23 options_path = os.path.join(this_path, 'input_files/OPTIONS-000005') 24 self.db_options = DatabaseOptions(options_path) 25 26 def test_setup(self): 27 self.assertEqual(self.bench_runner.db_bench_binary, self.pos_args[0]) 28 self.assertEqual(self.bench_runner.benchmark, self.pos_args[1]) 29 self.assertSetEqual( 30 set(self.bench_runner.db_bench_args), set(self.pos_args[2:]) 31 ) 32 33 def test_get_info_log_file_name(self): 34 log_file_name = DBBenchRunner.get_info_log_file_name( 35 None, 'random_path' 36 ) 37 self.assertEqual(log_file_name, 'LOG') 38 39 log_file_name = DBBenchRunner.get_info_log_file_name( 40 '/dev/shm/', '/tmp/rocksdbtest-155919/dbbench/' 41 ) 42 self.assertEqual(log_file_name, 'tmp_rocksdbtest-155919_dbbench_LOG') 43 44 def test_get_opt_args_str(self): 45 misc_opt_dict = {'bloom_bits': 2, 'empty_opt': None, 'rate_limiter': 3} 46 optional_args_str = DBBenchRunner.get_opt_args_str(misc_opt_dict) 47 self.assertEqual(optional_args_str, ' --bloom_bits=2 --rate_limiter=3') 48 49 def test_get_log_options(self): 50 db_path = '/tmp/rocksdb-155919/dbbench' 51 # when db_log_dir is present in the db_options 52 update_dict = { 53 'DBOptions.db_log_dir': {NO_COL_FAMILY: '/dev/shm'}, 54 'DBOptions.stats_dump_period_sec': {NO_COL_FAMILY: '20'} 55 } 56 self.db_options.update_options(update_dict) 57 log_file_prefix, stats_freq = self.bench_runner.get_log_options( 58 self.db_options, db_path 59 ) 60 self.assertEqual( 61 log_file_prefix, '/dev/shm/tmp_rocksdb-155919_dbbench_LOG' 62 ) 63 self.assertEqual(stats_freq, 20) 64 65 update_dict = { 66 'DBOptions.db_log_dir': {NO_COL_FAMILY: None}, 67 'DBOptions.stats_dump_period_sec': {NO_COL_FAMILY: '30'} 68 } 69 self.db_options.update_options(update_dict) 70 log_file_prefix, stats_freq = self.bench_runner.get_log_options( 71 self.db_options, db_path 72 ) 73 self.assertEqual(log_file_prefix, '/tmp/rocksdb-155919/dbbench/LOG') 74 self.assertEqual(stats_freq, 30) 75 76 def test_build_experiment_command(self): 77 # add some misc_options to db_options 78 update_dict = { 79 'bloom_bits': {NO_COL_FAMILY: 2}, 80 'rate_limiter_bytes_per_sec': {NO_COL_FAMILY: 128000000} 81 } 82 self.db_options.update_options(update_dict) 83 db_path = '/dev/shm' 84 experiment_command = self.bench_runner._build_experiment_command( 85 self.db_options, db_path 86 ) 87 opt_args_str = DBBenchRunner.get_opt_args_str( 88 self.db_options.get_misc_options() 89 ) 90 opt_args_str += ( 91 ' --options_file=' + 92 self.db_options.generate_options_config('12345') 93 ) 94 for arg in self.pos_args[2:]: 95 opt_args_str += (' --' + arg) 96 expected_command = ( 97 self.pos_args[0] + ' --benchmarks=' + self.pos_args[1] + 98 ' --statistics --perf_level=3 --db=' + db_path + opt_args_str 99 ) 100 self.assertEqual(experiment_command, expected_command) 101 102 103class TestDBBenchRunner(unittest.TestCase): 104 def setUp(self): 105 # Note: the db_bench binary should be present in the rocksdb/ directory 106 self.pos_args = [ 107 './../../db_bench', 108 'overwrite', 109 'use_existing_db=true', 110 'duration=20' 111 ] 112 self.bench_runner = DBBenchRunner(self.pos_args) 113 this_path = os.path.abspath(os.path.dirname(__file__)) 114 options_path = os.path.join(this_path, 'input_files/OPTIONS-000005') 115 self.db_options = DatabaseOptions(options_path) 116 117 def test_experiment_output(self): 118 update_dict = {'bloom_bits': {NO_COL_FAMILY: 2}} 119 self.db_options.update_options(update_dict) 120 db_path = '/dev/shm' 121 data_sources, throughput = self.bench_runner.run_experiment( 122 self.db_options, db_path 123 ) 124 self.assertEqual( 125 data_sources[DataSource.Type.DB_OPTIONS][0].type, 126 DataSource.Type.DB_OPTIONS 127 ) 128 self.assertEqual( 129 data_sources[DataSource.Type.LOG][0].type, 130 DataSource.Type.LOG 131 ) 132 self.assertEqual(len(data_sources[DataSource.Type.TIME_SERIES]), 2) 133 self.assertEqual( 134 data_sources[DataSource.Type.TIME_SERIES][0].type, 135 DataSource.Type.TIME_SERIES 136 ) 137 self.assertEqual( 138 data_sources[DataSource.Type.TIME_SERIES][1].type, 139 DataSource.Type.TIME_SERIES 140 ) 141 self.assertEqual( 142 data_sources[DataSource.Type.TIME_SERIES][1].stats_freq_sec, 0 143 ) 144 145 146if __name__ == '__main__': 147 unittest.main() 148