diff --git a/pkg/pool/sampler_test.go b/pkg/pool/sampler_test.go new file mode 100644 index 0000000..ed6a3f7 --- /dev/null +++ b/pkg/pool/sampler_test.go @@ -0,0 +1,40 @@ +package pool + +import ( + "math/rand" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestSamplerStability(t *testing.T) { + const COUNT = 100000 + + cases := []struct { + probabilities []float64 + expected []int + }{ + { + probabilities: []float64{1, 0}, + expected: []int{COUNT, 0}, + }, + { + probabilities: []float64{0.1, 0.2, 0.7}, + expected: []int{10138, 19813, 70049}, + }, + { + probabilities: []float64{0.2, 0.2, 0.4, 0.1, 0.1, 0}, + expected: []int{19824, 20169, 39900, 10243, 9864, 0}, + }, + } + + for _, tc := range cases { + sampler := NewSampler(tc.probabilities, rand.NewSource(0)) + res := make([]int, len(tc.probabilities)) + for i := 0; i < COUNT; i++ { + res[sampler.Next()]++ + } + + require.Equal(t, tc.expected, res, "probabilities: %v", tc.probabilities) + } +}