package pool import ( "context" "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) } } func TestHealthyReweight(t *testing.T) { var ( weights = []float64{0.9, 0.1} names = []string{"node0", "node1"} buffer = make([]float64, len(weights)) ) cache, err := newCache(0) require.NoError(t, err) client1 := newMockClient(names[0], *newPrivateKey(t)) client1.errOnDial() client2 := newMockClient(names[1], *newPrivateKey(t)) inner := &innerPool{ sampler: newSampler(weights, rand.NewSource(0)), clients: []client{client1, client2}, } p := &Pool{ innerPools: []*innerPool{inner}, cache: cache, key: newPrivateKey(t), rebalanceParams: rebalanceParameters{nodesParams: []*nodesParam{{weights: weights}}}, } // check getting first node connection before rebalance happened connection0, err := p.connection() require.NoError(t, err) mock0 := connection0.(*mockClient) require.Equal(t, names[0], mock0.address()) p.updateInnerNodesHealth(context.TODO(), 0, buffer) connection1, err := p.connection() require.NoError(t, err) mock1 := connection1.(*mockClient) require.Equal(t, names[1], mock1.address()) // enabled first node again inner.lock.Lock() inner.clients[0] = newMockClient(names[0], *newPrivateKey(t)) inner.lock.Unlock() p.updateInnerNodesHealth(context.TODO(), 0, buffer) inner.sampler = newSampler(weights, rand.NewSource(0)) connection0, err = p.connection() require.NoError(t, err) mock0 = connection0.(*mockClient) require.Equal(t, names[0], mock0.address()) } func TestHealthyNoReweight(t *testing.T) { var ( weights = []float64{0.9, 0.1} names = []string{"node0", "node1"} buffer = make([]float64, len(weights)) ) sampl := newSampler(weights, rand.NewSource(0)) inner := &innerPool{ sampler: sampl, clients: []client{ newMockClient(names[0], *newPrivateKey(t)), newMockClient(names[1], *newPrivateKey(t)), }, } p := &Pool{ innerPools: []*innerPool{inner}, rebalanceParams: rebalanceParameters{nodesParams: []*nodesParam{{weights: weights}}}, } p.updateInnerNodesHealth(context.TODO(), 0, buffer) inner.lock.RLock() defer inner.lock.RUnlock() require.Equal(t, inner.sampler, sampl) }