github.com/unionj-cloud/go-doudou@v1.3.8-0.20221011095552-0088008e5b31/framework/registry/serviceprovider_test.go (about) 1 package registry 2 3 import ( 4 "github.com/unionj-cloud/go-doudou/framework/memberlist" 5 "sync" 6 ) 7 8 type mockServiceProvider struct { 9 name string 10 servers []*memberlist.Node 11 serverMap map[string]*memberlist.Node 12 lock sync.RWMutex 13 } 14 15 func (m *mockServiceProvider) SelectServer() string { 16 return "mock server" 17 } 18 19 func (m *mockServiceProvider) AddNode(node *memberlist.Node) { 20 m.lock.Lock() 21 defer m.lock.Unlock() 22 m.servers = append(m.servers, node) 23 m.serverMap[node.Name] = node 24 } 25 26 func (m *mockServiceProvider) UpdateWeight(node *memberlist.Node) { 27 m.lock.Lock() 28 defer m.lock.Unlock() 29 if s, exists := m.serverMap[node.Name]; exists { 30 s.Weight = node.Weight 31 } 32 } 33 34 func (m *mockServiceProvider) RemoveNode(node *memberlist.Node) { 35 m.lock.Lock() 36 defer m.lock.Unlock() 37 if _, exists := m.serverMap[node.Name]; exists { 38 var idx int 39 for i, n := range m.servers { 40 if n.Name == node.Name { 41 idx = i 42 } 43 } 44 m.servers = append(m.servers[:idx], m.servers[idx+1:]...) 45 delete(m.serverMap, node.Name) 46 } 47 } 48 49 func newMockServiceProvider(name string) *mockServiceProvider { 50 return &mockServiceProvider{ 51 servers: make([]*memberlist.Node, 0), 52 serverMap: make(map[string]*memberlist.Node), 53 name: name, 54 } 55 }