1 // Copyright 2016 Google Inc. All rights reserved.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 // http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14
15 #include "src/weighted_reservoir_sampler.h"
16
17 #include <tuple>
18 #include <vector>
19
20 #include "port/gtest.h"
21
22 using testing::TestWithParam;
23 using testing::ValuesIn;
24 using testing::Combine;
25 using testing::Range;
26
27 namespace protobuf_mutator {
28
29 class WeightedReservoirSamplerTest
30 : public TestWithParam<std::tuple<int, std::vector<int>>> {};
31
32 const int kRuns = 1000000;
33
34 const std::vector<int> kTests[] = {
35 {1},
36 {1, 1, 1},
37 {1, 1, 0},
38 {1, 10, 100},
39 {100, 1, 10},
40 {1, 10000, 10000},
41 {1, 3, 7, 100, 105},
42 {93519, 52999, 354, 37837, 55285, 31787, 89096, 55695, 1587,
43 18233, 77557, 67632, 59348, 51250, 17417, 96856, 78568, 44296,
44 70170, 41328, 9206, 90187, 54086, 35602, 53167, 33791, 60118,
45 52962, 10327, 80513, 49526, 18326, 83662, 49644, 70903, 4910,
46 36309, 19196, 42982, 53316, 14773, 86607, 60835}};
47
48 INSTANTIATE_TEST_CASE_P(AllTest, WeightedReservoirSamplerTest,
49 Combine(Range(1, 10, 3), ValuesIn(kTests)));
50
TEST_P(WeightedReservoirSamplerTest,Test)51 TEST_P(WeightedReservoirSamplerTest, Test) {
52 std::vector<int> weights = std::get<1>(GetParam());
53 std::vector<int> counts(weights.size(), 0);
54
55 using RandomEngine = std::mt19937;
56 RandomEngine rand(std::get<0>(GetParam()));
57 for (int i = 0; i < kRuns; ++i) {
58 WeightedReservoirSampler<int, RandomEngine> sampler(&rand);
59 for (size_t j = 0; j < weights.size(); ++j) sampler.Try(weights[j], j);
60 ++counts[sampler.selected()];
61 }
62
63 int sum = std::accumulate(weights.begin(), weights.end(), 0);
64 for (size_t j = 0; j < weights.size(); ++j) {
65 float expected = weights[j];
66 expected /= sum;
67
68 float actual = counts[j];
69 actual /= kRuns;
70
71 EXPECT_NEAR(expected, actual, 0.01);
72 }
73 }
74
75 } // namespace protobuf_mutator
76