129 lines
4 KiB
Go
129 lines
4 KiB
Go
// Copyright (c) HashiCorp, Inc.
|
|
// SPDX-License-Identifier: MPL-2.0
|
|
|
|
package mock
|
|
|
|
import (
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/hashicorp/go-hclog"
|
|
"github.com/hashicorp/nomad/client/serviceregistration"
|
|
)
|
|
|
|
// Ensure that the mock handler implements the service registration handler
|
|
// interface.
|
|
var _ serviceregistration.Handler = (*ServiceRegistrationHandler)(nil)
|
|
|
|
// ServiceRegistrationHandler is the mock implementation of the
|
|
// serviceregistration.Handler interface and can be used for testing.
|
|
type ServiceRegistrationHandler struct {
|
|
log hclog.Logger
|
|
|
|
// ops tracks the requested operations by the caller during the entire
|
|
// lifecycle of the ServiceRegistrationHandler. The mutex should be used
|
|
// whenever interacting with this.
|
|
mu sync.Mutex
|
|
ops []Operation
|
|
|
|
// AllocRegistrationsFn allows injecting return values for the
|
|
// AllocRegistrations function.
|
|
AllocRegistrationsFn func(allocID string) (*serviceregistration.AllocRegistration, error)
|
|
}
|
|
|
|
// NewServiceRegistrationHandler returns a ready to use
|
|
// ServiceRegistrationHandler for testing.
|
|
func NewServiceRegistrationHandler(log hclog.Logger) *ServiceRegistrationHandler {
|
|
return &ServiceRegistrationHandler{
|
|
ops: make([]Operation, 0, 20),
|
|
log: log.Named("mock_service_registration"),
|
|
}
|
|
}
|
|
|
|
func (h *ServiceRegistrationHandler) RegisterWorkload(services *serviceregistration.WorkloadServices) error {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
|
|
h.log.Trace("RegisterWorkload", "alloc_id", services.AllocInfo.AllocID,
|
|
"name", services.Name(), "services", len(services.Services))
|
|
|
|
h.ops = append(h.ops, newOperation("add", services.AllocInfo.AllocID, services.Name()))
|
|
return nil
|
|
}
|
|
|
|
func (h *ServiceRegistrationHandler) RemoveWorkload(services *serviceregistration.WorkloadServices) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
|
|
h.log.Trace("RemoveWorkload", "alloc_id", services.AllocInfo.AllocID,
|
|
"name", services.Name(), "services", len(services.Services))
|
|
|
|
h.ops = append(h.ops, newOperation("remove", services.AllocInfo.AllocID, services.Name()))
|
|
}
|
|
|
|
func (h *ServiceRegistrationHandler) UpdateWorkload(old, newServices *serviceregistration.WorkloadServices) error {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
|
|
h.log.Trace("UpdateWorkload", "alloc_id", newServices.AllocInfo.AllocID, "name", newServices.Name(),
|
|
"old_services", len(old.Services), "new_services", len(newServices.Services))
|
|
|
|
h.ops = append(h.ops, newOperation("update", newServices.AllocInfo.AllocID, newServices.Name()))
|
|
return nil
|
|
}
|
|
|
|
func (h *ServiceRegistrationHandler) AllocRegistrations(allocID string) (*serviceregistration.AllocRegistration, error) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
|
|
h.log.Trace("AllocRegistrations", "alloc_id", allocID)
|
|
h.ops = append(h.ops, newOperation("alloc_registrations", allocID, ""))
|
|
|
|
if h.AllocRegistrationsFn != nil {
|
|
return h.AllocRegistrationsFn(allocID)
|
|
}
|
|
return nil, nil
|
|
}
|
|
|
|
func (h *ServiceRegistrationHandler) UpdateTTL(checkID, namespace, output, status string) error {
|
|
// TODO(tgross): this method is here so we can implement the
|
|
// interface but the locking we need for testing creates a lot
|
|
// of opportunities for deadlocks in testing that will never
|
|
// appear in live code.
|
|
h.log.Trace("UpdateTTL", "check_id", checkID, "namespace", namespace, "status", status)
|
|
return nil
|
|
}
|
|
|
|
// GetOps returns all stored operations within the handler.
|
|
func (h *ServiceRegistrationHandler) GetOps() []Operation {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
|
|
return h.ops
|
|
}
|
|
|
|
// Operation represents the register/deregister operations.
|
|
type Operation struct {
|
|
Op string // add, remove, or update
|
|
AllocID string
|
|
Name string // task or group name
|
|
OccurredAt time.Time
|
|
}
|
|
|
|
// newOperation generates a new Operation for the given parameters.
|
|
func newOperation(op, allocID, name string) Operation {
|
|
switch op {
|
|
case "add", "remove", "update", "alloc_registrations",
|
|
"add_group", "remove_group", "update_group", "update_ttl":
|
|
default:
|
|
panic(fmt.Errorf("invalid consul op: %s", op))
|
|
}
|
|
return Operation{
|
|
Op: op,
|
|
AllocID: allocID,
|
|
Name: name,
|
|
OccurredAt: time.Now(),
|
|
}
|
|
}
|