open-nomad/client/consul/consul_testing.go

package consul

import (
	"fmt"
	"sync"
	"time"

	log "github.com/hashicorp/go-hclog"
	"github.com/hashicorp/nomad/command/agent/consul"
	testing "github.com/mitchellh/go-testing-interface"
)

// MockConsulOp represents the register/deregister operations.
type MockConsulOp struct {
	Op         string // add, remove, or update
	AllocID    string
	Name       string // task or group name
	OccurredAt time.Time
}

func NewMockConsulOp(op, allocID, name string) MockConsulOp {
	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 MockConsulOp{
		Op:         op,
		AllocID:    allocID,
		Name:       name,
		OccurredAt: time.Now(),
	}
}

// MockConsulServiceClient implements the ConsulServiceAPI interface to record
// and log task registration/deregistration.
type MockConsulServiceClient struct {
	ops []MockConsulOp
	mu  sync.Mutex

	logger log.Logger

	// AllocRegistrationsFn allows injecting return values for the
	// AllocRegistrations function.
	AllocRegistrationsFn func(allocID string) (*consul.AllocRegistration, error)
}

func NewMockConsulServiceClient(t testing.T, logger log.Logger) *MockConsulServiceClient {
	logger = logger.Named("mock_consul")
	m := MockConsulServiceClient{
		ops:    make([]MockConsulOp, 0, 20),
		logger: logger,
	}
	return &m
}

func (m *MockConsulServiceClient) UpdateWorkload(old, newSvcs *consul.WorkloadServices) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.logger.Trace("UpdateWorkload", "alloc_id", newSvcs.AllocID, "name", newSvcs.Name(),
		"old_services", len(old.Services), "new_services", len(newSvcs.Services),
	)
	m.ops = append(m.ops, NewMockConsulOp("update", newSvcs.AllocID, newSvcs.Name()))
	return nil
}

func (m *MockConsulServiceClient) RegisterWorkload(svcs *consul.WorkloadServices) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.logger.Trace("RegisterWorkload", "alloc_id", svcs.AllocID, "name", svcs.Name(),
		"services", len(svcs.Services),
	)
	m.ops = append(m.ops, NewMockConsulOp("add", svcs.AllocID, svcs.Name()))
	return nil
}

func (m *MockConsulServiceClient) RemoveWorkload(svcs *consul.WorkloadServices) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.logger.Trace("RemoveWorkload", "alloc_id", svcs.AllocID, "name", svcs.Name(),
		"services", len(svcs.Services),
	)
	m.ops = append(m.ops, NewMockConsulOp("remove", svcs.AllocID, svcs.Name()))
}

func (m *MockConsulServiceClient) AllocRegistrations(allocID string) (*consul.AllocRegistration, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.logger.Trace("AllocRegistrations", "alloc_id", allocID)
	m.ops = append(m.ops, NewMockConsulOp("alloc_registrations", allocID, ""))

	if m.AllocRegistrationsFn != nil {
		return m.AllocRegistrationsFn(allocID)
	}

	return nil, nil
}

func (m *MockConsulServiceClient) UpdateTTL(checkID, 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.
	m.logger.Trace("UpdateTTL", "check_id", checkID, "status", status)
	return nil
}

func (m *MockConsulServiceClient) GetOps() []MockConsulOp {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.ops
}