From eeeaa8a4fd8f981f843668103ee2ffa7a15b2276 Mon Sep 17 00:00:00 2001
From: Denis Kirillov <denis@nspcc.ru>
Date: Fri, 18 Jun 2021 14:09:43 +0300
Subject: [PATCH] [#14] Added tests for sampler

Signed-off-by: Denis Kirillov <denis@nspcc.ru>
---
 pkg/pool/sampler_test.go | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)
 create mode 100644 pkg/pool/sampler_test.go

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)
+	}
+}