github.com/blixtra/nomad@v0.7.2-0.20171221000451-da9a1d7bb050/client/consul_test.go (about) 1 package client 2 3 import ( 4 "fmt" 5 "io/ioutil" 6 "log" 7 "os" 8 "sync" 9 "testing" 10 11 "github.com/hashicorp/nomad/client/driver" 12 cstructs "github.com/hashicorp/nomad/client/structs" 13 "github.com/hashicorp/nomad/command/agent/consul" 14 "github.com/hashicorp/nomad/nomad/structs" 15 ) 16 17 // mockConsulOp represents the register/deregister operations. 18 type mockConsulOp struct { 19 op string // add, remove, or update 20 allocID string 21 task *structs.Task 22 exec driver.ScriptExecutor 23 net *cstructs.DriverNetwork 24 } 25 26 func newMockConsulOp(op, allocID string, task *structs.Task, exec driver.ScriptExecutor, net *cstructs.DriverNetwork) mockConsulOp { 27 if op != "add" && op != "remove" && op != "update" && op != "alloc_registrations" { 28 panic(fmt.Errorf("invalid consul op: %s", op)) 29 } 30 return mockConsulOp{ 31 op: op, 32 allocID: allocID, 33 task: task, 34 exec: exec, 35 net: net, 36 } 37 } 38 39 // mockConsulServiceClient implements the ConsulServiceAPI interface to record 40 // and log task registration/deregistration. 41 type mockConsulServiceClient struct { 42 ops []mockConsulOp 43 mu sync.Mutex 44 45 logger *log.Logger 46 47 // allocRegistrationsFn allows injecting return values for the 48 // AllocRegistrations function. 49 allocRegistrationsFn func(allocID string) (*consul.AllocRegistration, error) 50 } 51 52 func newMockConsulServiceClient() *mockConsulServiceClient { 53 m := mockConsulServiceClient{ 54 ops: make([]mockConsulOp, 0, 20), 55 logger: log.New(ioutil.Discard, "", 0), 56 } 57 if testing.Verbose() { 58 m.logger = log.New(os.Stderr, "", log.LstdFlags) 59 } 60 return &m 61 } 62 63 func (m *mockConsulServiceClient) UpdateTask(allocID string, old, new *structs.Task, restarter consul.TaskRestarter, exec driver.ScriptExecutor, net *cstructs.DriverNetwork) error { 64 m.mu.Lock() 65 defer m.mu.Unlock() 66 m.logger.Printf("[TEST] mock_consul: UpdateTask(%q, %v, %v, %T, %x)", allocID, old, new, exec, net.Hash()) 67 m.ops = append(m.ops, newMockConsulOp("update", allocID, new, exec, net)) 68 return nil 69 } 70 71 func (m *mockConsulServiceClient) RegisterTask(allocID string, task *structs.Task, restarter consul.TaskRestarter, exec driver.ScriptExecutor, net *cstructs.DriverNetwork) error { 72 m.mu.Lock() 73 defer m.mu.Unlock() 74 m.logger.Printf("[TEST] mock_consul: RegisterTask(%q, %q, %T, %x)", allocID, task.Name, exec, net.Hash()) 75 m.ops = append(m.ops, newMockConsulOp("add", allocID, task, exec, net)) 76 return nil 77 } 78 79 func (m *mockConsulServiceClient) RemoveTask(allocID string, task *structs.Task) { 80 m.mu.Lock() 81 defer m.mu.Unlock() 82 m.logger.Printf("[TEST] mock_consul: RemoveTask(%q, %q)", allocID, task.Name) 83 m.ops = append(m.ops, newMockConsulOp("remove", allocID, task, nil, nil)) 84 } 85 86 func (m *mockConsulServiceClient) AllocRegistrations(allocID string) (*consul.AllocRegistration, error) { 87 m.mu.Lock() 88 defer m.mu.Unlock() 89 m.logger.Printf("[TEST] mock_consul: AllocRegistrations(%q)", allocID) 90 m.ops = append(m.ops, newMockConsulOp("alloc_registrations", allocID, nil, nil, nil)) 91 92 if m.allocRegistrationsFn != nil { 93 return m.allocRegistrationsFn(allocID) 94 } 95 96 return nil, nil 97 }