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  }