github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/client/serviceregistration/mock/mock.go (about)

     1  package mock
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/hashicorp/go-hclog"
     9  	"github.com/hashicorp/nomad/client/serviceregistration"
    10  )
    11  
    12  // Ensure that the mock handler implements the service registration handler
    13  // interface.
    14  var _ serviceregistration.Handler = (*ServiceRegistrationHandler)(nil)
    15  
    16  // ServiceRegistrationHandler is the mock implementation of the
    17  // serviceregistration.Handler interface and can be used for testing.
    18  type ServiceRegistrationHandler struct {
    19  	log hclog.Logger
    20  
    21  	// ops tracks the requested operations by the caller during the entire
    22  	// lifecycle of the ServiceRegistrationHandler. The mutex should be used
    23  	// whenever interacting with this.
    24  	mu  sync.Mutex
    25  	ops []Operation
    26  
    27  	// AllocRegistrationsFn allows injecting return values for the
    28  	// AllocRegistrations function.
    29  	AllocRegistrationsFn func(allocID string) (*serviceregistration.AllocRegistration, error)
    30  }
    31  
    32  // NewServiceRegistrationHandler returns a ready to use
    33  // ServiceRegistrationHandler for testing.
    34  func NewServiceRegistrationHandler(log hclog.Logger) *ServiceRegistrationHandler {
    35  	return &ServiceRegistrationHandler{
    36  		ops: make([]Operation, 0, 20),
    37  		log: log.Named("mock_service_registration"),
    38  	}
    39  }
    40  
    41  func (h *ServiceRegistrationHandler) RegisterWorkload(services *serviceregistration.WorkloadServices) error {
    42  	h.mu.Lock()
    43  	defer h.mu.Unlock()
    44  
    45  	h.log.Trace("RegisterWorkload", "alloc_id", services.AllocInfo.AllocID,
    46  		"name", services.Name(), "services", len(services.Services))
    47  
    48  	h.ops = append(h.ops, newOperation("add", services.AllocInfo.AllocID, services.Name()))
    49  	return nil
    50  }
    51  
    52  func (h *ServiceRegistrationHandler) RemoveWorkload(services *serviceregistration.WorkloadServices) {
    53  	h.mu.Lock()
    54  	defer h.mu.Unlock()
    55  
    56  	h.log.Trace("RemoveWorkload", "alloc_id", services.AllocInfo.AllocID,
    57  		"name", services.Name(), "services", len(services.Services))
    58  
    59  	h.ops = append(h.ops, newOperation("remove", services.AllocInfo.AllocID, services.Name()))
    60  }
    61  
    62  func (h *ServiceRegistrationHandler) UpdateWorkload(old, newServices *serviceregistration.WorkloadServices) error {
    63  	h.mu.Lock()
    64  	defer h.mu.Unlock()
    65  
    66  	h.log.Trace("UpdateWorkload", "alloc_id", newServices.AllocInfo.AllocID, "name", newServices.Name(),
    67  		"old_services", len(old.Services), "new_services", len(newServices.Services))
    68  
    69  	h.ops = append(h.ops, newOperation("update", newServices.AllocInfo.AllocID, newServices.Name()))
    70  	return nil
    71  }
    72  
    73  func (h *ServiceRegistrationHandler) AllocRegistrations(allocID string) (*serviceregistration.AllocRegistration, error) {
    74  	h.mu.Lock()
    75  	defer h.mu.Unlock()
    76  
    77  	h.log.Trace("AllocRegistrations", "alloc_id", allocID)
    78  	h.ops = append(h.ops, newOperation("alloc_registrations", allocID, ""))
    79  
    80  	if h.AllocRegistrationsFn != nil {
    81  		return h.AllocRegistrationsFn(allocID)
    82  	}
    83  	return nil, nil
    84  }
    85  
    86  func (h *ServiceRegistrationHandler) UpdateTTL(checkID, namespace, output, status string) error {
    87  	// TODO(tgross): this method is here so we can implement the
    88  	// interface but the locking we need for testing creates a lot
    89  	// of opportunities for deadlocks in testing that will never
    90  	// appear in live code.
    91  	h.log.Trace("UpdateTTL", "check_id", checkID, "namespace", namespace, "status", status)
    92  	return nil
    93  }
    94  
    95  // GetOps returns all stored operations within the handler.
    96  func (h *ServiceRegistrationHandler) GetOps() []Operation {
    97  	h.mu.Lock()
    98  	defer h.mu.Unlock()
    99  
   100  	return h.ops
   101  }
   102  
   103  // Operation represents the register/deregister operations.
   104  type Operation struct {
   105  	Op         string // add, remove, or update
   106  	AllocID    string
   107  	Name       string // task or group name
   108  	OccurredAt time.Time
   109  }
   110  
   111  // newOperation generates a new Operation for the given parameters.
   112  func newOperation(op, allocID, name string) Operation {
   113  	switch op {
   114  	case "add", "remove", "update", "alloc_registrations",
   115  		"add_group", "remove_group", "update_group", "update_ttl":
   116  	default:
   117  		panic(fmt.Errorf("invalid consul op: %s", op))
   118  	}
   119  	return Operation{
   120  		Op:         op,
   121  		AllocID:    allocID,
   122  		Name:       name,
   123  		OccurredAt: time.Now(),
   124  	}
   125  }