github.com/blixtra/nomad@v0.7.2-0.20171221000451-da9a1d7bb050/command/agent/consul/catalog_testing.go (about)

     1  package consul
     2  
     3  import (
     4  	"fmt"
     5  	"log"
     6  	"sync"
     7  
     8  	"github.com/hashicorp/consul/api"
     9  )
    10  
    11  // MockCatalog can be used for testing where the CatalogAPI is needed.
    12  type MockCatalog struct {
    13  	logger *log.Logger
    14  }
    15  
    16  func NewMockCatalog(l *log.Logger) *MockCatalog {
    17  	return &MockCatalog{logger: l}
    18  }
    19  
    20  func (m *MockCatalog) Datacenters() ([]string, error) {
    21  	dcs := []string{"dc1"}
    22  	m.logger.Printf("[DEBUG] mock_consul: Datacenters() -> (%q, nil)", dcs)
    23  	return dcs, nil
    24  }
    25  
    26  func (m *MockCatalog) Service(service, tag string, q *api.QueryOptions) ([]*api.CatalogService, *api.QueryMeta, error) {
    27  	m.logger.Printf("[DEBUG] mock_consul: Service(%q, %q, %#v) -> (nil, nil, nil)", service, tag, q)
    28  	return nil, nil, nil
    29  }
    30  
    31  // MockAgent is a fake in-memory Consul backend for ServiceClient.
    32  type MockAgent struct {
    33  	// maps of what services and checks have been registered
    34  	services map[string]*api.AgentServiceRegistration
    35  	checks   map[string]*api.AgentCheckRegistration
    36  	mu       sync.Mutex
    37  
    38  	// when UpdateTTL is called the check ID will have its counter inc'd
    39  	checkTTLs map[string]int
    40  
    41  	// What check status to return from Checks()
    42  	checkStatus string
    43  }
    44  
    45  // NewMockAgent that returns all checks as passing.
    46  func NewMockAgent() *MockAgent {
    47  	return &MockAgent{
    48  		services:    make(map[string]*api.AgentServiceRegistration),
    49  		checks:      make(map[string]*api.AgentCheckRegistration),
    50  		checkTTLs:   make(map[string]int),
    51  		checkStatus: api.HealthPassing,
    52  	}
    53  }
    54  
    55  // SetStatus that Checks() should return. Returns old status value.
    56  func (c *MockAgent) SetStatus(s string) string {
    57  	c.mu.Lock()
    58  	old := c.checkStatus
    59  	c.checkStatus = s
    60  	c.mu.Unlock()
    61  	return old
    62  }
    63  
    64  func (c *MockAgent) Services() (map[string]*api.AgentService, error) {
    65  	c.mu.Lock()
    66  	defer c.mu.Unlock()
    67  
    68  	r := make(map[string]*api.AgentService, len(c.services))
    69  	for k, v := range c.services {
    70  		r[k] = &api.AgentService{
    71  			ID:                v.ID,
    72  			Service:           v.Name,
    73  			Tags:              make([]string, len(v.Tags)),
    74  			Port:              v.Port,
    75  			Address:           v.Address,
    76  			EnableTagOverride: v.EnableTagOverride,
    77  		}
    78  		copy(r[k].Tags, v.Tags)
    79  	}
    80  	return r, nil
    81  }
    82  
    83  // Checks implements the Agent API Checks method.
    84  func (c *MockAgent) Checks() (map[string]*api.AgentCheck, error) {
    85  	c.mu.Lock()
    86  	defer c.mu.Unlock()
    87  
    88  	r := make(map[string]*api.AgentCheck, len(c.checks))
    89  	for k, v := range c.checks {
    90  		r[k] = &api.AgentCheck{
    91  			CheckID:     v.ID,
    92  			Name:        v.Name,
    93  			Status:      c.checkStatus,
    94  			Notes:       v.Notes,
    95  			ServiceID:   v.ServiceID,
    96  			ServiceName: c.services[v.ServiceID].Name,
    97  		}
    98  	}
    99  	return r, nil
   100  }
   101  
   102  // CheckRegs returns the raw AgentCheckRegistrations registered with this mock
   103  // agent.
   104  func (c *MockAgent) CheckRegs() []*api.AgentCheckRegistration {
   105  	c.mu.Lock()
   106  	defer c.mu.Unlock()
   107  
   108  	regs := make([]*api.AgentCheckRegistration, 0, len(c.checks))
   109  	for _, check := range c.checks {
   110  		regs = append(regs, check)
   111  	}
   112  	return regs
   113  }
   114  
   115  func (c *MockAgent) CheckRegister(check *api.AgentCheckRegistration) error {
   116  	c.mu.Lock()
   117  	defer c.mu.Unlock()
   118  	c.checks[check.ID] = check
   119  
   120  	// Be nice and make checks reachable-by-service
   121  	scheck := check.AgentServiceCheck
   122  	c.services[check.ServiceID].Checks = append(c.services[check.ServiceID].Checks, &scheck)
   123  	return nil
   124  }
   125  
   126  func (c *MockAgent) CheckDeregister(checkID string) error {
   127  	c.mu.Lock()
   128  	defer c.mu.Unlock()
   129  	delete(c.checks, checkID)
   130  	delete(c.checkTTLs, checkID)
   131  	return nil
   132  }
   133  
   134  func (c *MockAgent) ServiceRegister(service *api.AgentServiceRegistration) error {
   135  	c.mu.Lock()
   136  	defer c.mu.Unlock()
   137  	c.services[service.ID] = service
   138  	return nil
   139  }
   140  
   141  func (c *MockAgent) ServiceDeregister(serviceID string) error {
   142  	c.mu.Lock()
   143  	defer c.mu.Unlock()
   144  	delete(c.services, serviceID)
   145  	return nil
   146  }
   147  
   148  func (c *MockAgent) UpdateTTL(id string, output string, status string) error {
   149  	c.mu.Lock()
   150  	defer c.mu.Unlock()
   151  	check, ok := c.checks[id]
   152  	if !ok {
   153  		return fmt.Errorf("unknown check id: %q", id)
   154  	}
   155  	// Flip initial status to passing
   156  	check.Status = "passing"
   157  	c.checkTTLs[id]++
   158  	return nil
   159  }