github.com/zhizhiboom/nomad@v0.8.5-0.20180907175415-f28fd3a1a056/client/consul/consul_testing.go (about)

     1  package consul
     2  
     3  import (
     4  	"fmt"
     5  	"log"
     6  	"sync"
     7  
     8  	"github.com/hashicorp/nomad/command/agent/consul"
     9  	"github.com/hashicorp/nomad/helper/testlog"
    10  	"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) *MockConsulServiceClient {
    45  	m := MockConsulServiceClient{
    46  		Ops:    make([]MockConsulOp, 0, 20),
    47  		Logger: testlog.Logger(t),
    48  	}
    49  	return &m
    50  }
    51  
    52  func (m *MockConsulServiceClient) UpdateTask(old, new *consul.TaskServices) error {
    53  	m.mu.Lock()
    54  	defer m.mu.Unlock()
    55  	m.Logger.Printf("[TEST] mock_consul: UpdateTask(alloc: %s, task: %s)", new.AllocID[:6], new.Name)
    56  	m.Ops = append(m.Ops, NewMockConsulOp("update", new.AllocID, new.Name))
    57  	return nil
    58  }
    59  
    60  func (m *MockConsulServiceClient) RegisterTask(task *consul.TaskServices) error {
    61  	m.mu.Lock()
    62  	defer m.mu.Unlock()
    63  	m.Logger.Printf("[TEST] mock_consul: RegisterTask(alloc: %s, task: %s)", task.AllocID, task.Name)
    64  	m.Ops = append(m.Ops, NewMockConsulOp("add", task.AllocID, task.Name))
    65  	return nil
    66  }
    67  
    68  func (m *MockConsulServiceClient) RemoveTask(task *consul.TaskServices) {
    69  	m.mu.Lock()
    70  	defer m.mu.Unlock()
    71  	m.Logger.Printf("[TEST] mock_consul: RemoveTask(%q, %q)", task.AllocID, task.Name)
    72  	m.Ops = append(m.Ops, NewMockConsulOp("remove", task.AllocID, task.Name))
    73  }
    74  
    75  func (m *MockConsulServiceClient) AllocRegistrations(allocID string) (*consul.AllocRegistration, error) {
    76  	m.mu.Lock()
    77  	defer m.mu.Unlock()
    78  	m.Logger.Printf("[TEST] mock_consul: AllocRegistrations(%q)", allocID)
    79  	m.Ops = append(m.Ops, NewMockConsulOp("alloc_registrations", allocID, ""))
    80  
    81  	if m.AllocRegistrationsFn != nil {
    82  		return m.AllocRegistrationsFn(allocID)
    83  	}
    84  
    85  	return nil, nil
    86  }