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