github.com/iqoqo/nomad@v0.11.3-0.20200911112621-d7021c74d101/nomad/drainer/drain_testing.go (about)

     1  package drainer
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/hashicorp/nomad/nomad/structs"
     7  )
     8  
     9  type MockNodeTrackerEvent struct {
    10  	NodeUpdate *structs.Node
    11  	NodeRemove string
    12  }
    13  
    14  type MockNodeTracker struct {
    15  	Nodes  map[string]*structs.Node
    16  	Events []*MockNodeTrackerEvent
    17  	sync.Mutex
    18  }
    19  
    20  func NewMockNodeTracker() *MockNodeTracker {
    21  	return &MockNodeTracker{
    22  		Nodes:  make(map[string]*structs.Node),
    23  		Events: make([]*MockNodeTrackerEvent, 0, 16),
    24  	}
    25  }
    26  
    27  func (m *MockNodeTracker) TrackedNodes() map[string]*structs.Node {
    28  	m.Lock()
    29  	defer m.Unlock()
    30  	return m.Nodes
    31  }
    32  
    33  func (m *MockNodeTracker) Remove(nodeID string) {
    34  	m.Lock()
    35  	defer m.Unlock()
    36  	delete(m.Nodes, nodeID)
    37  	m.Events = append(m.Events, &MockNodeTrackerEvent{NodeRemove: nodeID})
    38  }
    39  
    40  func (m *MockNodeTracker) Update(node *structs.Node) {
    41  	m.Lock()
    42  	defer m.Unlock()
    43  	m.Nodes[node.ID] = node
    44  	m.Events = append(m.Events, &MockNodeTrackerEvent{NodeUpdate: node})
    45  }
    46  
    47  func (m *MockNodeTracker) events() []*MockNodeTrackerEvent {
    48  	m.Lock()
    49  	defer m.Unlock()
    50  
    51  	return m.Events
    52  }