github.com/manicqin/nomad@v0.9.5/client/consul/consul_testing.go (about)

     1  package consul
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  	"time"
     7  
     8  	log "github.com/hashicorp/go-hclog"
     9  
    10  	"github.com/hashicorp/nomad/command/agent/consul"
    11  	testing "github.com/mitchellh/go-testing-interface"
    12  )
    13  
    14  // MockConsulOp represents the register/deregister operations.
    15  type MockConsulOp struct {
    16  	Op         string // add, remove, or update
    17  	AllocID    string
    18  	Name       string // task or group name
    19  	OccurredAt time.Time
    20  }
    21  
    22  func NewMockConsulOp(op, allocID, name string) MockConsulOp {
    23  	switch op {
    24  	case "add", "remove", "update", "alloc_registrations",
    25  		"add_group", "remove_group", "update_group", "update_ttl":
    26  	default:
    27  		panic(fmt.Errorf("invalid consul op: %s", op))
    28  	}
    29  	return MockConsulOp{
    30  		Op:         op,
    31  		AllocID:    allocID,
    32  		Name:       name,
    33  		OccurredAt: time.Now(),
    34  	}
    35  }
    36  
    37  // MockConsulServiceClient implements the ConsulServiceAPI interface to record
    38  // and log task registration/deregistration.
    39  type MockConsulServiceClient struct {
    40  	ops []MockConsulOp
    41  	mu  sync.Mutex
    42  
    43  	logger log.Logger
    44  
    45  	// AllocRegistrationsFn allows injecting return values for the
    46  	// AllocRegistrations function.
    47  	AllocRegistrationsFn func(allocID string) (*consul.AllocRegistration, error)
    48  }
    49  
    50  func NewMockConsulServiceClient(t testing.T, logger log.Logger) *MockConsulServiceClient {
    51  	logger = logger.Named("mock_consul")
    52  	m := MockConsulServiceClient{
    53  		ops:    make([]MockConsulOp, 0, 20),
    54  		logger: logger,
    55  	}
    56  	return &m
    57  }
    58  
    59  func (m *MockConsulServiceClient) UpdateWorkload(old, newSvcs *consul.WorkloadServices) error {
    60  	m.mu.Lock()
    61  	defer m.mu.Unlock()
    62  	m.logger.Trace("UpdateWorkload", "alloc_id", newSvcs.AllocID, "name", newSvcs.Name(),
    63  		"old_services", len(old.Services), "new_services", len(newSvcs.Services),
    64  	)
    65  	m.ops = append(m.ops, NewMockConsulOp("update", newSvcs.AllocID, newSvcs.Name()))
    66  	return nil
    67  }
    68  
    69  func (m *MockConsulServiceClient) RegisterWorkload(svcs *consul.WorkloadServices) error {
    70  	m.mu.Lock()
    71  	defer m.mu.Unlock()
    72  	m.logger.Trace("RegisterWorkload", "alloc_id", svcs.AllocID, "name", svcs.Name(),
    73  		"services", len(svcs.Services),
    74  	)
    75  	m.ops = append(m.ops, NewMockConsulOp("add", svcs.AllocID, svcs.Name()))
    76  	return nil
    77  }
    78  
    79  func (m *MockConsulServiceClient) RemoveWorkload(svcs *consul.WorkloadServices) {
    80  	m.mu.Lock()
    81  	defer m.mu.Unlock()
    82  	m.logger.Trace("RemoveWorkload", "alloc_id", svcs.AllocID, "name", svcs.Name(),
    83  		"services", len(svcs.Services),
    84  	)
    85  	m.ops = append(m.ops, NewMockConsulOp("remove", svcs.AllocID, svcs.Name()))
    86  }
    87  
    88  func (m *MockConsulServiceClient) AllocRegistrations(allocID string) (*consul.AllocRegistration, error) {
    89  	m.mu.Lock()
    90  	defer m.mu.Unlock()
    91  	m.logger.Trace("AllocRegistrations", "alloc_id", allocID)
    92  	m.ops = append(m.ops, NewMockConsulOp("alloc_registrations", allocID, ""))
    93  
    94  	if m.AllocRegistrationsFn != nil {
    95  		return m.AllocRegistrationsFn(allocID)
    96  	}
    97  
    98  	return nil, nil
    99  }
   100  
   101  func (m *MockConsulServiceClient) UpdateTTL(checkID, output, status string) error {
   102  	// TODO(tgross): this method is here so we can implement the
   103  	// interface but the locking we need for testing creates a lot
   104  	// of opportunities for deadlocks in testing that will never
   105  	// appear in live code.
   106  	m.logger.Trace("UpdateTTL", "check_id", checkID, "status", status)
   107  	return nil
   108  }
   109  
   110  func (m *MockConsulServiceClient) GetOps() []MockConsulOp {
   111  	m.mu.Lock()
   112  	defer m.mu.Unlock()
   113  	return m.ops
   114  }