github.com/smintz/nomad@v0.8.3/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  
    37  	// hits is the total number of times agent methods have been called
    38  	hits int
    39  
    40  	// mu guards above fields
    41  	mu sync.Mutex
    42  
    43  	// when UpdateTTL is called the check ID will have its counter inc'd
    44  	checkTTLs map[string]int
    45  
    46  	// What check status to return from Checks()
    47  	checkStatus string
    48  }
    49  
    50  // NewMockAgent that returns all checks as passing.
    51  func NewMockAgent() *MockAgent {
    52  	return &MockAgent{
    53  		services:    make(map[string]*api.AgentServiceRegistration),
    54  		checks:      make(map[string]*api.AgentCheckRegistration),
    55  		checkTTLs:   make(map[string]int),
    56  		checkStatus: api.HealthPassing,
    57  	}
    58  }
    59  
    60  // getHits returns how many Consul Agent API calls have been made.
    61  func (c *MockAgent) getHits() int {
    62  	c.mu.Lock()
    63  	defer c.mu.Unlock()
    64  	return c.hits
    65  }
    66  
    67  // SetStatus that Checks() should return. Returns old status value.
    68  func (c *MockAgent) SetStatus(s string) string {
    69  	c.mu.Lock()
    70  	old := c.checkStatus
    71  	c.checkStatus = s
    72  	c.mu.Unlock()
    73  	return old
    74  }
    75  
    76  func (c *MockAgent) Self() (map[string]map[string]interface{}, error) {
    77  	c.mu.Lock()
    78  	defer c.mu.Unlock()
    79  	c.hits++
    80  
    81  	s := map[string]map[string]interface{}{
    82  		"Member": {
    83  			"Addr":        "127.0.0.1",
    84  			"DelegateCur": 4,
    85  			"DelegateMax": 5,
    86  			"DelegateMin": 2,
    87  			"Name":        "rusty",
    88  			"Port":        8301,
    89  			"ProtocolCur": 2,
    90  			"ProtocolMax": 5,
    91  			"ProtocolMin": 1,
    92  			"Status":      1,
    93  			"Tags": map[string]interface{}{
    94  				"build": "0.8.1:'e9ca44d",
    95  			},
    96  		},
    97  	}
    98  	return s, nil
    99  }
   100  
   101  func (c *MockAgent) Services() (map[string]*api.AgentService, error) {
   102  	c.mu.Lock()
   103  	defer c.mu.Unlock()
   104  	c.hits++
   105  
   106  	r := make(map[string]*api.AgentService, len(c.services))
   107  	for k, v := range c.services {
   108  		r[k] = &api.AgentService{
   109  			ID:                v.ID,
   110  			Service:           v.Name,
   111  			Tags:              make([]string, len(v.Tags)),
   112  			Port:              v.Port,
   113  			Address:           v.Address,
   114  			EnableTagOverride: v.EnableTagOverride,
   115  		}
   116  		copy(r[k].Tags, v.Tags)
   117  	}
   118  	return r, nil
   119  }
   120  
   121  // Checks implements the Agent API Checks method.
   122  func (c *MockAgent) Checks() (map[string]*api.AgentCheck, error) {
   123  	c.mu.Lock()
   124  	defer c.mu.Unlock()
   125  	c.hits++
   126  
   127  	r := make(map[string]*api.AgentCheck, len(c.checks))
   128  	for k, v := range c.checks {
   129  		r[k] = &api.AgentCheck{
   130  			CheckID:     v.ID,
   131  			Name:        v.Name,
   132  			Status:      c.checkStatus,
   133  			Notes:       v.Notes,
   134  			ServiceID:   v.ServiceID,
   135  			ServiceName: c.services[v.ServiceID].Name,
   136  		}
   137  	}
   138  	return r, nil
   139  }
   140  
   141  // CheckRegs returns the raw AgentCheckRegistrations registered with this mock
   142  // agent.
   143  func (c *MockAgent) CheckRegs() []*api.AgentCheckRegistration {
   144  	c.mu.Lock()
   145  	defer c.mu.Unlock()
   146  	regs := make([]*api.AgentCheckRegistration, 0, len(c.checks))
   147  	for _, check := range c.checks {
   148  		regs = append(regs, check)
   149  	}
   150  	return regs
   151  }
   152  
   153  func (c *MockAgent) CheckRegister(check *api.AgentCheckRegistration) error {
   154  	c.mu.Lock()
   155  	defer c.mu.Unlock()
   156  	c.hits++
   157  
   158  	c.checks[check.ID] = check
   159  
   160  	// Be nice and make checks reachable-by-service
   161  	scheck := check.AgentServiceCheck
   162  	c.services[check.ServiceID].Checks = append(c.services[check.ServiceID].Checks, &scheck)
   163  	return nil
   164  }
   165  
   166  func (c *MockAgent) CheckDeregister(checkID string) error {
   167  	c.mu.Lock()
   168  	defer c.mu.Unlock()
   169  	c.hits++
   170  	delete(c.checks, checkID)
   171  	delete(c.checkTTLs, checkID)
   172  	return nil
   173  }
   174  
   175  func (c *MockAgent) ServiceRegister(service *api.AgentServiceRegistration) error {
   176  	c.mu.Lock()
   177  	defer c.mu.Unlock()
   178  	c.hits++
   179  	c.services[service.ID] = service
   180  	return nil
   181  }
   182  
   183  func (c *MockAgent) ServiceDeregister(serviceID string) error {
   184  	c.mu.Lock()
   185  	defer c.mu.Unlock()
   186  	c.hits++
   187  	delete(c.services, serviceID)
   188  	return nil
   189  }
   190  
   191  func (c *MockAgent) UpdateTTL(id string, output string, status string) error {
   192  	c.mu.Lock()
   193  	defer c.mu.Unlock()
   194  	c.hits++
   195  
   196  	check, ok := c.checks[id]
   197  	if !ok {
   198  		return fmt.Errorf("unknown check id: %q", id)
   199  	}
   200  	// Flip initial status to passing
   201  	check.Status = "passing"
   202  	c.checkTTLs[id]++
   203  	return nil
   204  }