-
Notifications
You must be signed in to change notification settings - Fork 0
/
picker_test.go
78 lines (68 loc) · 1.79 KB
/
picker_test.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
package goutube
import (
"testing"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/balancer/base"
"google.golang.org/grpc/resolver"
"github.com/stretchr/testify/require"
)
func TestPickerNoSubConnAvailable(t *testing.T) {
picker := &Picker{}
for _, method := range []string{
"/goutube.vX.goutube/Produce",
"/goutube.vX.goutube/Consume",
} {
info := balancer.PickInfo{
FullMethodName: method,
}
result, err := picker.Pick(info)
require.Equal(t, balancer.ErrNoSubConnAvailable, err)
require.Nil(t, result.SubConn)
}
}
func TestPickerProducesToLeader(t *testing.T) {
picker, subConns := setupTestPicker()
info := balancer.PickInfo{
FullMethodName: "/goutube.vX.goutube/Produce",
}
for i := 0; i < 5; i++ {
gotPick, err := picker.Pick(info)
require.NoError(t, err)
require.Equal(t, subConns[0], gotPick.SubConn)
}
}
func TestPickerConsumesFromFollowers(t *testing.T) {
picker, subConns := setupTestPicker()
info := balancer.PickInfo{
FullMethodName: "/goutube.vX.goutube/Consume",
}
for i := 0; i < 5; i++ {
pick, err := picker.Pick(info)
require.NoError(t, err)
require.Equal(t, subConns[i%2+1], pick.SubConn)
}
}
func setupTestPicker() (*Picker, []*subConn) {
var subConns []*subConn
buildInfo := base.PickerBuildInfo{
ReadySCs: make(map[balancer.SubConn]base.SubConnInfo),
}
for i := 0; i < 3; i++ {
sc := &subConn{}
addr := resolver.Address{}
// 0th sub conn is the leader
sc.UpdateAddresses([]resolver.Address{addr})
buildInfo.ReadySCs[sc] = base.SubConnInfo{Address: addr}
subConns = append(subConns, sc)
}
picker := &Picker{}
picker.Build(buildInfo)
return picker, subConns
}
type subConn struct {
addrs []resolver.Address
}
func (s *subConn) UpdateAddresses(addrs []resolver.Address) {
s.addrs = addrs
}
func (s *subConn) Connect() {}