github.com/lmb/consul@v1.4.1/testutil/server_methods.go (about) 1 package testutil 2 3 import ( 4 "bytes" 5 "encoding/base64" 6 "encoding/json" 7 "fmt" 8 "io" 9 "io/ioutil" 10 "log" 11 "net/http" 12 "testing" 13 14 "github.com/pkg/errors" 15 ) 16 17 // copied from testutil to break circular dependency 18 const ( 19 HealthAny = "any" 20 HealthPassing = "passing" 21 HealthWarning = "warning" 22 HealthCritical = "critical" 23 HealthMaint = "maintenance" 24 ) 25 26 // JoinLAN is used to join local datacenters together. 27 func (s *TestServer) JoinLAN(t *testing.T, addr string) { 28 resp := s.put(t, "/v1/agent/join/"+addr, nil) 29 defer resp.Body.Close() 30 } 31 32 // JoinWAN is used to join remote datacenters together. 33 func (s *TestServer) JoinWAN(t *testing.T, addr string) { 34 resp := s.put(t, "/v1/agent/join/"+addr+"?wan=1", nil) 35 resp.Body.Close() 36 } 37 38 // SetKV sets an individual key in the K/V store. 39 func (s *TestServer) SetKV(t *testing.T, key string, val []byte) { 40 resp := s.put(t, "/v1/kv/"+key, bytes.NewBuffer(val)) 41 resp.Body.Close() 42 } 43 44 // SetKVString sets an individual key in the K/V store, but accepts a string 45 // instead of []byte. 46 func (s *TestServer) SetKVString(t *testing.T, key string, val string) { 47 resp := s.put(t, "/v1/kv/"+key, bytes.NewBufferString(val)) 48 resp.Body.Close() 49 } 50 51 // GetKV retrieves a single key and returns its value 52 func (s *TestServer) GetKV(t *testing.T, key string) []byte { 53 resp := s.get(t, "/v1/kv/"+key) 54 defer resp.Body.Close() 55 56 raw, err := ioutil.ReadAll(resp.Body) 57 if err != nil { 58 t.Fatalf("failed to read body: %s", err) 59 } 60 61 var result []*TestKVResponse 62 if err := json.Unmarshal(raw, &result); err != nil { 63 t.Fatalf("failed to unmarshal: %s", err) 64 } 65 if len(result) < 1 { 66 t.Fatalf("key does not exist: %s", key) 67 } 68 69 v, err := base64.StdEncoding.DecodeString(result[0].Value) 70 if err != nil { 71 t.Fatalf("failed to base64 decode: %s", err) 72 } 73 74 return v 75 } 76 77 // GetKVString retrieves a value from the store, but returns as a string instead 78 // of []byte. 79 func (s *TestServer) GetKVString(t *testing.T, key string) string { 80 return string(s.GetKV(t, key)) 81 } 82 83 // PopulateKV fills the Consul KV with data from a generic map. 84 func (s *TestServer) PopulateKV(t *testing.T, data map[string][]byte) { 85 for k, v := range data { 86 s.SetKV(t, k, v) 87 } 88 } 89 90 // ListKV returns a list of keys present in the KV store. This will list all 91 // keys under the given prefix recursively and return them as a slice. 92 func (s *TestServer) ListKV(t *testing.T, prefix string) []string { 93 resp := s.get(t, "/v1/kv/"+prefix+"?keys") 94 defer resp.Body.Close() 95 96 raw, err := ioutil.ReadAll(resp.Body) 97 if err != nil { 98 t.Fatalf("failed to read body: %s", err) 99 } 100 101 var result []string 102 if err := json.Unmarshal(raw, &result); err != nil { 103 t.Fatalf("failed to unmarshal: %s", err) 104 } 105 return result 106 } 107 108 // AddService adds a new service to the Consul instance. It also 109 // automatically adds a health check with the given status, which 110 // can be one of "passing", "warning", or "critical". 111 func (s *TestServer) AddService(t *testing.T, name, status string, tags []string) { 112 s.AddAddressableService(t, name, status, "", 0, tags) // set empty address and 0 as port for non-accessible service 113 } 114 115 // AddAddressableService adds a new service to the Consul instance by 116 // passing "address" and "port". It is helpful when you need to prepare a fakeService 117 // that maybe accessed with in target source code. 118 // It also automatically adds a health check with the given status, which 119 // can be one of "passing", "warning", or "critical", just like `AddService` does. 120 func (s *TestServer) AddAddressableService(t *testing.T, name, status, address string, port int, tags []string) { 121 svc := &TestService{ 122 Name: name, 123 Tags: tags, 124 Address: address, 125 Port: port, 126 } 127 payload, err := s.encodePayload(svc) 128 if err != nil { 129 t.Fatal(err) 130 } 131 s.put(t, "/v1/agent/service/register", payload) 132 133 chkName := "service:" + name 134 chk := &TestCheck{ 135 Name: chkName, 136 ServiceID: name, 137 TTL: "10m", 138 } 139 payload, err = s.encodePayload(chk) 140 if err != nil { 141 t.Fatal(err) 142 } 143 s.put(t, "/v1/agent/check/register", payload) 144 145 switch status { 146 case HealthPassing: 147 s.put(t, "/v1/agent/check/pass/"+chkName, nil) 148 case HealthWarning: 149 s.put(t, "/v1/agent/check/warn/"+chkName, nil) 150 case HealthCritical: 151 s.put(t, "/v1/agent/check/fail/"+chkName, nil) 152 default: 153 t.Fatalf("Unrecognized status: %s", status) 154 } 155 } 156 157 // AddCheck adds a check to the Consul instance. If the serviceID is 158 // left empty (""), then the check will be associated with the node. 159 // The check status may be "passing", "warning", or "critical". 160 func (s *TestServer) AddCheck(t *testing.T, name, serviceID, status string) { 161 chk := &TestCheck{ 162 ID: name, 163 Name: name, 164 TTL: "10m", 165 } 166 if serviceID != "" { 167 chk.ServiceID = serviceID 168 } 169 170 payload, err := s.encodePayload(chk) 171 if err != nil { 172 t.Fatal(err) 173 } 174 s.put(t, "/v1/agent/check/register", payload) 175 176 switch status { 177 case HealthPassing: 178 s.put(t, "/v1/agent/check/pass/"+name, nil) 179 case HealthWarning: 180 s.put(t, "/v1/agent/check/warn/"+name, nil) 181 case HealthCritical: 182 s.put(t, "/v1/agent/check/fail/"+name, nil) 183 default: 184 t.Fatalf("Unrecognized status: %s", status) 185 } 186 } 187 188 // put performs a new HTTP PUT request. 189 func (s *TestServer) put(t *testing.T, path string, body io.Reader) *http.Response { 190 req, err := http.NewRequest("PUT", s.url(path), body) 191 if err != nil { 192 t.Fatalf("failed to create PUT request: %s", err) 193 } 194 resp, err := s.HTTPClient.Do(req) 195 if err != nil { 196 t.Fatalf("failed to make PUT request: %s", err) 197 } 198 if err := s.requireOK(resp); err != nil { 199 defer resp.Body.Close() 200 t.Fatalf("not OK PUT: %s", err) 201 } 202 return resp 203 } 204 205 // get performs a new HTTP GET request. 206 func (s *TestServer) get(t *testing.T, path string) *http.Response { 207 resp, err := s.HTTPClient.Get(s.url(path)) 208 if err != nil { 209 t.Fatalf("failed to create GET request: %s", err) 210 } 211 if err := s.requireOK(resp); err != nil { 212 defer resp.Body.Close() 213 t.Fatalf("not OK GET: %s", err) 214 } 215 return resp 216 } 217 218 // encodePayload returns a new io.Reader wrapping the encoded contents 219 // of the payload, suitable for passing directly to a new request. 220 func (s *TestServer) encodePayload(payload interface{}) (io.Reader, error) { 221 var encoded bytes.Buffer 222 enc := json.NewEncoder(&encoded) 223 if err := enc.Encode(payload); err != nil { 224 return nil, errors.Wrap(err, "failed to encode payload") 225 } 226 return &encoded, nil 227 } 228 229 // url is a helper function which takes a relative URL and 230 // makes it into a proper URL against the local Consul server. 231 func (s *TestServer) url(path string) string { 232 if s == nil { 233 log.Fatal("s is nil") 234 } 235 if s.Config == nil { 236 log.Fatal("s.Config is nil") 237 } 238 if s.Config.Ports == nil { 239 log.Fatal("s.Config.Ports is nil") 240 } 241 if s.Config.Ports.HTTP == 0 { 242 log.Fatal("s.Config.Ports.HTTP is 0") 243 } 244 if path == "" { 245 log.Fatal("path is empty") 246 } 247 return fmt.Sprintf("http://127.0.0.1:%d%s", s.Config.Ports.HTTP, path) 248 } 249 250 // requireOK checks the HTTP response code and ensures it is acceptable. 251 func (s *TestServer) requireOK(resp *http.Response) error { 252 if resp.StatusCode != 200 { 253 return fmt.Errorf("Bad status code: %d", resp.StatusCode) 254 } 255 return nil 256 }