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  }