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 }