frostfs-sdk-go/pool/sampler_test.go

142 lines
3.5 KiB
Go
Raw Normal View History

package pool
import (
"context"
"fmt"
"math/rand"
"testing"
"github.com/nspcc-dev/neofs-sdk-go/client"
"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)
}
}
type clientMock struct {
client.Client
name string
err error
}
func (c *clientMock) EndpointInfo(context.Context, ...client.CallOption) (*client.EndpointInfoRes, error) {
return nil, nil
}
func (c *clientMock) NetworkInfo(context.Context, ...client.CallOption) (*client.NetworkInfoRes, error) {
return nil, nil
}
func newNetmapMock(name string, needErr bool) *clientMock {
var err error
if needErr {
err = fmt.Errorf("not available")
}
return &clientMock{name: name, err: err}
}
func TestHealthyReweight(t *testing.T) {
var (
weights = []float64{0.9, 0.1}
names = []string{"node0", "node1"}
options = &BuilderOptions{nodesParams: []*NodesParam{{weights: weights}}}
buffer = make([]float64, len(weights))
)
cache, err := NewCache()
require.NoError(t, err)
inner := &innerPool{
sampler: NewSampler(weights, rand.NewSource(0)),
clientPacks: []*clientPack{
{client: newNetmapMock(names[0], true), healthy: true, address: "address0"},
{client: newNetmapMock(names[1], false), healthy: true, address: "address1"}},
}
p := &pool{
innerPools: []*innerPool{inner},
cache: cache,
key: newPrivateKey(t),
}
// check getting first node connection before rebalance happened
connection0, _, err := p.Connection()
require.NoError(t, err)
mock0 := connection0.(*clientMock)
require.Equal(t, names[0], mock0.name)
updateInnerNodesHealth(context.TODO(), p, 0, options, buffer)
connection1, _, err := p.Connection()
require.NoError(t, err)
mock1 := connection1.(*clientMock)
require.Equal(t, names[1], mock1.name)
// enabled first node again
inner.lock.Lock()
inner.clientPacks[0].client = newNetmapMock(names[0], false)
inner.lock.Unlock()
updateInnerNodesHealth(context.TODO(), p, 0, options, buffer)
inner.sampler = NewSampler(weights, rand.NewSource(0))
connection0, _, err = p.Connection()
require.NoError(t, err)
mock0 = connection0.(*clientMock)
require.Equal(t, names[0], mock0.name)
}
func TestHealthyNoReweight(t *testing.T) {
var (
weights = []float64{0.9, 0.1}
names = []string{"node0", "node1"}
options = &BuilderOptions{nodesParams: []*NodesParam{{weights: weights}}}
buffer = make([]float64, len(weights))
)
sampler := NewSampler(weights, rand.NewSource(0))
inner := &innerPool{
sampler: sampler,
clientPacks: []*clientPack{
{client: newNetmapMock(names[0], false), healthy: true},
{client: newNetmapMock(names[1], false), healthy: true}},
}
p := &pool{
innerPools: []*innerPool{inner},
}
updateInnerNodesHealth(context.TODO(), p, 0, options, buffer)
inner.lock.RLock()
defer inner.lock.RUnlock()
require.Equal(t, inner.sampler, sampler)
}