github.com/smithx10/nomad@v0.9.1-rc1/client/consul/consul_testing.go (about) 1 package consul 2 3 import ( 4 "fmt" 5 "sync" 6 7 log "github.com/hashicorp/go-hclog" 8 9 "github.com/hashicorp/nomad/command/agent/consul" 10 testing "github.com/mitchellh/go-testing-interface" 11 ) 12 13 // MockConsulOp represents the register/deregister operations. 14 type MockConsulOp struct { 15 Op string // add, remove, or update 16 AllocID string 17 Task string 18 } 19 20 func NewMockConsulOp(op, allocID, task string) MockConsulOp { 21 if op != "add" && op != "remove" && op != "update" && op != "alloc_registrations" { 22 panic(fmt.Errorf("invalid consul op: %s", op)) 23 } 24 return MockConsulOp{ 25 Op: op, 26 AllocID: allocID, 27 Task: task, 28 } 29 } 30 31 // MockConsulServiceClient implements the ConsulServiceAPI interface to record 32 // and log task registration/deregistration. 33 type MockConsulServiceClient struct { 34 ops []MockConsulOp 35 mu sync.Mutex 36 37 logger log.Logger 38 39 // AllocRegistrationsFn allows injecting return values for the 40 // AllocRegistrations function. 41 AllocRegistrationsFn func(allocID string) (*consul.AllocRegistration, error) 42 } 43 44 func NewMockConsulServiceClient(t testing.T, logger log.Logger) *MockConsulServiceClient { 45 logger = logger.Named("mock_consul") 46 m := MockConsulServiceClient{ 47 ops: make([]MockConsulOp, 0, 20), 48 logger: logger, 49 } 50 return &m 51 } 52 53 func (m *MockConsulServiceClient) UpdateTask(old, newSvcs *consul.TaskServices) error { 54 m.mu.Lock() 55 defer m.mu.Unlock() 56 m.logger.Trace("UpdateTask", "alloc_id", newSvcs.AllocID, "task", newSvcs.Name, 57 "old_services", len(old.Services), "new_services", len(newSvcs.Services), 58 ) 59 m.ops = append(m.ops, NewMockConsulOp("update", newSvcs.AllocID, newSvcs.Name)) 60 return nil 61 } 62 63 func (m *MockConsulServiceClient) RegisterTask(task *consul.TaskServices) error { 64 m.mu.Lock() 65 defer m.mu.Unlock() 66 m.logger.Trace("RegisterTask", "alloc_id", task.AllocID, "task", task.Name, 67 "services", len(task.Services), 68 ) 69 m.ops = append(m.ops, NewMockConsulOp("add", task.AllocID, task.Name)) 70 return nil 71 } 72 73 func (m *MockConsulServiceClient) RemoveTask(task *consul.TaskServices) { 74 m.mu.Lock() 75 defer m.mu.Unlock() 76 m.logger.Trace("RemoveTask", "alloc_id", task.AllocID, "task", task.Name, 77 "services", len(task.Services), 78 ) 79 m.ops = append(m.ops, NewMockConsulOp("remove", task.AllocID, task.Name)) 80 } 81 82 func (m *MockConsulServiceClient) AllocRegistrations(allocID string) (*consul.AllocRegistration, error) { 83 m.mu.Lock() 84 defer m.mu.Unlock() 85 m.logger.Trace("AllocRegistrations", "alloc_id", allocID) 86 m.ops = append(m.ops, NewMockConsulOp("alloc_registrations", allocID, "")) 87 88 if m.AllocRegistrationsFn != nil { 89 return m.AllocRegistrationsFn(allocID) 90 } 91 92 return nil, nil 93 } 94 95 func (m *MockConsulServiceClient) GetOps() []MockConsulOp { 96 m.mu.Lock() 97 defer m.mu.Unlock() 98 return m.ops 99 }