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

package consul

import (
	"fmt"
	"sync"

	"github.com/hashicorp/consul/api"
	"github.com/hashicorp/go-hclog"
	"github.com/hashicorp/nomad/helper"
)

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

func NewMockCatalog(l hclog.Logger) *MockCatalog {
	return &MockCatalog{logger: l.Named("mock_consul")}
}

func (m *MockCatalog) Datacenters() ([]string, error) {
	dcs := []string{"dc1"}
	m.logger.Trace("Datacenters()", "dcs", dcs, "error", "nil")
	return dcs, nil
}

func (m *MockCatalog) Service(service, tag string, q *api.QueryOptions) ([]*api.CatalogService, *api.QueryMeta, error) {
	m.logger.Trace("Services()", "service", service, "tag", tag, "query_options", 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

	// hits is the total number of times agent methods have been called
	hits int

	// mu guards above fields
	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,
	}
}

// getHits returns how many Consul Agent API calls have been made.
func (c *MockAgent) getHits() int {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.hits
}

// 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) Self() (map[string]map[string]interface{}, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.hits++

	s := map[string]map[string]interface{}{
		"Member": {
			"Addr":        "127.0.0.1",
			"DelegateCur": 4,
			"DelegateMax": 5,
			"DelegateMin": 2,
			"Name":        "rusty",
			"Port":        8301,
			"ProtocolCur": 2,
			"ProtocolMax": 5,
			"ProtocolMin": 1,
			"Status":      1,
			"Tags": map[string]interface{}{
				"build": "0.8.1:'e9ca44d",
			},
		},
	}
	return s, nil
}

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

	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)),
			Meta:              helper.CopyMapStringString(v.Meta),
			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()
	c.hits++

	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.hits++

	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()
	c.hits++
	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.hits++
	c.services[service.ID] = service
	return nil
}

func (c *MockAgent) ServiceDeregister(serviceID string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.hits++
	delete(c.services, serviceID)
	for k, v := range c.checks {
		if v.ServiceID == serviceID {
			delete(c.checks, k)
			delete(c.checkTTLs, k)
		}
	}
	return nil
}

func (c *MockAgent) UpdateTTL(id string, output string, status string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.hits++

	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
}

// a convenience method for looking up a registered service by name
func (c *MockAgent) lookupService(name string) []*api.AgentServiceRegistration {
	c.mu.Lock()
	defer c.mu.Unlock()

	var services []*api.AgentServiceRegistration
	for _, service := range c.services {
		if service.Name == name {
			services = append(services, service)
		}
	}
	return services
}