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 <numeric>
18 #include <tuple>
19 #include <vector>
20
21 #include "port/gtest.h"
22
23 using testing::Combine;
24 using testing::Range;
25 using testing::TestWithParam;
26 using testing::ValuesIn;
27
28 namespace protobuf_mutator {
29
30 class WeightedReservoirSamplerTest
31 : public TestWithParam<std::tuple<int, std::vector<int>>> {};
32
33 const int kRuns = 1000000;
34
35 const std::vector<int> kTests[] = {
36 {1},
37 {1, 1, 1},
38 {1, 1, 0},
39 {1, 10, 100},
40 {100, 1, 10},
41 {1, 10000, 10000},
42 {1, 3, 7, 100, 105},
43 {93519, 52999, 354, 37837, 55285, 31787, 89096, 55695, 1587,
44 18233, 77557, 67632, 59348, 51250, 17417, 96856, 78568, 44296,
45 70170, 41328, 9206, 90187, 54086, 35602, 53167, 33791, 60118,
46 52962, 10327, 80513, 49526, 18326, 83662, 49644, 70903, 4910,
47 36309, 19196, 42982, 53316, 14773, 86607, 60835}};
48
49 INSTANTIATE_TEST_SUITE_P(AllTest, WeightedReservoirSamplerTest,
50 Combine(Range(1, 10, 3), ValuesIn(kTests)));
51
TEST_P(WeightedReservoirSamplerTest,Test)52 TEST_P(WeightedReservoirSamplerTest, Test) {
53 std::vector<int> weights = std::get<1>(GetParam());
54 std::vector<int> counts(weights.size(), 0);
55
56 using RandomEngine = std::minstd_rand;
57 RandomEngine rand(std::get<0>(GetParam()));
58 for (int i = 0; i < kRuns; ++i) {
59 WeightedReservoirSampler<int, RandomEngine> sampler(&rand);
60 for (size_t j = 0; j < weights.size(); ++j) sampler.Try(weights[j], j);
61 ++counts[sampler.selected()];
62 }
63
64 int sum = std::accumulate(weights.begin(), weights.end(), 0);
65 for (size_t j = 0; j < weights.size(); ++j) {
66 float expected = weights[j];
67 expected /= sum;
68
69 float actual = counts[j];
70 actual /= kRuns;
71
72 EXPECT_NEAR(expected, actual, 0.01);
73 }
74 }
75
76 } // namespace protobuf_mutator
77