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