open-nomad/command/agent/consul/catalog_testing.go

package consul

import (
	"fmt"
	"log"
	"sync"

	"github.com/hashicorp/consul/api"
)

// MockCatalog can be used for testing where the CatalogAPI is needed.
type MockCatalog struct {
	logger *log.Logger
}

func NewMockCatalog(l *log.Logger) *MockCatalog {
	return &MockCatalog{logger: l}
}

func (m *MockCatalog) Datacenters() ([]string, error) {
	dcs := []string{"dc1"}
	m.logger.Printf("[DEBUG] mock_consul: Datacenters() -> (%q, nil)", dcs)
	return dcs, nil
}

func (m *MockCatalog) Service(service, tag string, q *api.QueryOptions) ([]*api.CatalogService, *api.QueryMeta, error) {
	m.logger.Printf("[DEBUG] mock_consul: Service(%q, %q, %#v) -> (nil, nil, nil)", service, tag, q)
	return nil, nil, nil
}

// MockAgent is a fake in-memory Consul backend for ServiceClient.
type MockAgent struct {
	// maps of what services and checks have been registered
	services map[string]*api.AgentServiceRegistration
	checks   map[string]*api.AgentCheckRegistration
	mu       sync.Mutex

	// when UpdateTTL is called the check ID will have its counter inc'd
	checkTTLs map[string]int

	// What check status to return from Checks()
	checkStatus string
}

// NewMockAgent that returns all checks as passing.
func NewMockAgent() *MockAgent {
	return &MockAgent{
		services:    make(map[string]*api.AgentServiceRegistration),
		checks:      make(map[string]*api.AgentCheckRegistration),
		checkTTLs:   make(map[string]int),
		checkStatus: api.HealthPassing,
	}
}

// SetStatus that Checks() should return. Returns old status value.
func (c *MockAgent) SetStatus(s string) string {
	c.mu.Lock()
	old := c.checkStatus
	c.checkStatus = s
	c.mu.Unlock()
	return old
}

func (c *MockAgent) Services() (map[string]*api.AgentService, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	r := make(map[string]*api.AgentService, len(c.services))
	for k, v := range c.services {
		r[k] = &api.AgentService{
			ID:                v.ID,
			Service:           v.Name,
			Tags:              make([]string, len(v.Tags)),
			Port:              v.Port,
			Address:           v.Address,
			EnableTagOverride: v.EnableTagOverride,
		}
		copy(r[k].Tags, v.Tags)
	}
	return r, nil
}

// Checks implements the Agent API Checks method.
func (c *MockAgent) Checks() (map[string]*api.AgentCheck, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	r := make(map[string]*api.AgentCheck, len(c.checks))
	for k, v := range c.checks {
		r[k] = &api.AgentCheck{
			CheckID:     v.ID,
			Name:        v.Name,
			Status:      c.checkStatus,
			Notes:       v.Notes,
			ServiceID:   v.ServiceID,
			ServiceName: c.services[v.ServiceID].Name,
		}
	}
	return r, nil
}

// CheckRegs returns the raw AgentCheckRegistrations registered with this mock
// agent.
func (c *MockAgent) CheckRegs() []*api.AgentCheckRegistration {
	c.mu.Lock()
	defer c.mu.Unlock()

	regs := make([]*api.AgentCheckRegistration, 0, len(c.checks))
	for _, check := range c.checks {
		regs = append(regs, check)
	}
	return regs
}

func (c *MockAgent) CheckRegister(check *api.AgentCheckRegistration) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.checks[check.ID] = check

	// Be nice and make checks reachable-by-service
	scheck := check.AgentServiceCheck
	c.services[check.ServiceID].Checks = append(c.services[check.ServiceID].Checks, &scheck)
	return nil
}

func (c *MockAgent) CheckDeregister(checkID string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.checks, checkID)
	delete(c.checkTTLs, checkID)
	return nil
}

func (c *MockAgent) ServiceRegister(service *api.AgentServiceRegistration) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.services[service.ID] = service
	return nil
}

func (c *MockAgent) ServiceDeregister(serviceID string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.services, serviceID)
	return nil
}

func (c *MockAgent) UpdateTTL(id string, output string, status string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	check, ok := c.checks[id]
	if !ok {
		return fmt.Errorf("unknown check id: %q", id)
	}
	// Flip initial status to passing
	check.Status = "passing"
	c.checkTTLs[id]++
	return nil
}