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