github.com/huiliang/nomad@v0.2.1-0.20151124023127-7a8b664699ff/testutil/server.go (about) 1 package testutil 2 3 // TestServer is a test helper. It uses a fork/exec model to create 4 // a test Nomad server instance in the background and initialize it 5 // with some data and/or services. The test server can then be used 6 // to run a unit test, and offers an easy API to tear itself down 7 // when the test has completed. The only prerequisite is to have a nomad 8 // binary available on the $PATH. 9 // 10 // This package does not use Nomad's official API client. This is 11 // because we use TestServer to test the API client, which would 12 // otherwise cause an import cycle. 13 14 import ( 15 "bytes" 16 "encoding/json" 17 "fmt" 18 "io" 19 "io/ioutil" 20 "net/http" 21 "os" 22 "os/exec" 23 "sync/atomic" 24 "testing" 25 26 "github.com/hashicorp/go-cleanhttp" 27 ) 28 29 // offset is used to atomically increment the port numbers. 30 var offset uint64 31 32 // TestServerConfig is the main server configuration struct. 33 type TestServerConfig struct { 34 NodeName string `json:"name,omitempty"` 35 DataDir string `json:"data_dir,omitempty"` 36 Region string `json:"region,omitempty"` 37 DisableCheckpoint bool `json:"disable_update_check"` 38 LogLevel string `json:"log_level,omitempty"` 39 Ports *PortsConfig `json:"ports,omitempty"` 40 Server *ServerConfig `json:"server,omitempty"` 41 Client *ClientConfig `json:"client,omitempty"` 42 DevMode bool `json:"-"` 43 Stdout, Stderr io.Writer `json:"-"` 44 } 45 46 // Ports is used to configure the network ports we use. 47 type PortsConfig struct { 48 HTTP int `json:"http,omitempty"` 49 RPC int `json:"rpc,omitempty"` 50 Serf int `json:"serf,omitempty"` 51 } 52 53 // ServerConfig is used to configure the nomad server. 54 type ServerConfig struct { 55 Enabled bool `json:"enabled"` 56 BootstrapExpect int `json:"bootstrap_expect"` 57 } 58 59 // ClientConfig is used to configure the client 60 type ClientConfig struct { 61 Enabled bool `json:"enabled"` 62 } 63 64 // ServerConfigCallback is a function interface which can be 65 // passed to NewTestServerConfig to modify the server config. 66 type ServerConfigCallback func(c *TestServerConfig) 67 68 // defaultServerConfig returns a new TestServerConfig struct 69 // with all of the listen ports incremented by one. 70 func defaultServerConfig() *TestServerConfig { 71 idx := int(atomic.AddUint64(&offset, 1)) 72 73 return &TestServerConfig{ 74 NodeName: fmt.Sprintf("node%d", idx), 75 DisableCheckpoint: true, 76 LogLevel: "DEBUG", 77 Ports: &PortsConfig{ 78 HTTP: 20000 + idx, 79 RPC: 21000 + idx, 80 Serf: 22000 + idx, 81 }, 82 Server: &ServerConfig{ 83 Enabled: true, 84 BootstrapExpect: 1, 85 }, 86 Client: &ClientConfig{ 87 Enabled: false, 88 }, 89 } 90 } 91 92 // TestServer is the main server wrapper struct. 93 type TestServer struct { 94 PID int 95 Config *TestServerConfig 96 t *testing.T 97 98 HTTPAddr string 99 SerfAddr string 100 HttpClient *http.Client 101 } 102 103 // NewTestServerConfig creates a new TestServer, and makes a call to 104 // an optional callback function to modify the configuration. 105 func NewTestServer(t *testing.T, cb ServerConfigCallback) *TestServer { 106 if path, err := exec.LookPath("nomad"); err != nil || path == "" { 107 t.Skip("nomad not found on $PATH, skipping") 108 } 109 110 dataDir, err := ioutil.TempDir("", "nomad") 111 if err != nil { 112 t.Fatalf("err: %s", err) 113 } 114 115 configFile, err := ioutil.TempFile(dataDir, "nomad") 116 if err != nil { 117 defer os.RemoveAll(dataDir) 118 t.Fatalf("err: %s", err) 119 } 120 121 nomadConfig := defaultServerConfig() 122 nomadConfig.DataDir = dataDir 123 124 if cb != nil { 125 cb(nomadConfig) 126 } 127 128 configContent, err := json.Marshal(nomadConfig) 129 if err != nil { 130 t.Fatalf("err: %s", err) 131 } 132 133 if _, err := configFile.Write(configContent); err != nil { 134 t.Fatalf("err: %s", err) 135 } 136 configFile.Close() 137 138 stdout := io.Writer(os.Stdout) 139 if nomadConfig.Stdout != nil { 140 stdout = nomadConfig.Stdout 141 } 142 143 stderr := io.Writer(os.Stderr) 144 if nomadConfig.Stderr != nil { 145 stderr = nomadConfig.Stderr 146 } 147 148 args := []string{"agent", "-config", configFile.Name()} 149 if nomadConfig.DevMode { 150 args = append(args, "-dev") 151 } 152 153 // Start the server 154 cmd := exec.Command("nomad", args...) 155 cmd.Stdout = stdout 156 cmd.Stderr = stderr 157 if err := cmd.Start(); err != nil { 158 t.Fatalf("err: %s", err) 159 } 160 161 client := cleanhttp.DefaultClient() 162 163 server := &TestServer{ 164 Config: nomadConfig, 165 PID: cmd.Process.Pid, 166 t: t, 167 168 HTTPAddr: fmt.Sprintf("127.0.0.1:%d", nomadConfig.Ports.HTTP), 169 SerfAddr: fmt.Sprintf("127.0.0.1:%d", nomadConfig.Ports.Serf), 170 HttpClient: client, 171 } 172 173 // Wait for the server to be ready 174 if nomadConfig.Server.Enabled && nomadConfig.Server.BootstrapExpect != 0 { 175 server.waitForLeader() 176 } else { 177 server.waitForAPI() 178 } 179 return server 180 } 181 182 // Stop stops the test Nomad server, and removes the Nomad data 183 // directory once we are done. 184 func (s *TestServer) Stop() { 185 defer os.RemoveAll(s.Config.DataDir) 186 187 cmd := exec.Command("kill", "-9", fmt.Sprintf("%d", s.PID)) 188 if err := cmd.Run(); err != nil { 189 s.t.Errorf("err: %s", err) 190 } 191 } 192 193 // waitForAPI waits for only the agent HTTP endpoint to start 194 // responding. This is an indication that the agent has started, 195 // but will likely return before a leader is elected. 196 func (s *TestServer) waitForAPI() { 197 WaitForResult(func() (bool, error) { 198 resp, err := s.HttpClient.Get(s.url("/v1/agent/self")) 199 if err != nil { 200 return false, err 201 } 202 defer resp.Body.Close() 203 if err := s.requireOK(resp); err != nil { 204 return false, err 205 } 206 return true, nil 207 }, func(err error) { 208 defer s.Stop() 209 s.t.Fatalf("err: %s", err) 210 }) 211 } 212 213 // waitForLeader waits for the Nomad server's HTTP API to become 214 // available, and then waits for a known leader and an index of 215 // 1 or more to be observed to confirm leader election is done. 216 func (s *TestServer) waitForLeader() { 217 WaitForResult(func() (bool, error) { 218 // Query the API and check the status code 219 resp, err := s.HttpClient.Get(s.url("/v1/jobs")) 220 if err != nil { 221 return false, err 222 } 223 defer resp.Body.Close() 224 if err := s.requireOK(resp); err != nil { 225 return false, err 226 } 227 228 // Ensure we have a leader and a node registeration 229 if leader := resp.Header.Get("X-Nomad-KnownLeader"); leader != "true" { 230 return false, fmt.Errorf("Nomad leader status: %#v", leader) 231 } 232 return true, nil 233 }, func(err error) { 234 defer s.Stop() 235 s.t.Fatalf("err: %s", err) 236 }) 237 } 238 239 // url is a helper function which takes a relative URL and 240 // makes it into a proper URL against the local Nomad server. 241 func (s *TestServer) url(path string) string { 242 return fmt.Sprintf("http://%s%s", s.HTTPAddr, path) 243 } 244 245 // requireOK checks the HTTP response code and ensures it is acceptable. 246 func (s *TestServer) requireOK(resp *http.Response) error { 247 if resp.StatusCode != 200 { 248 return fmt.Errorf("Bad status code: %d", resp.StatusCode) 249 } 250 return nil 251 } 252 253 // put performs a new HTTP PUT request. 254 func (s *TestServer) put(path string, body io.Reader) *http.Response { 255 req, err := http.NewRequest("PUT", s.url(path), body) 256 if err != nil { 257 s.t.Fatalf("err: %s", err) 258 } 259 resp, err := s.HttpClient.Do(req) 260 if err != nil { 261 s.t.Fatalf("err: %s", err) 262 } 263 if err := s.requireOK(resp); err != nil { 264 defer resp.Body.Close() 265 s.t.Fatal(err) 266 } 267 return resp 268 } 269 270 // get performs a new HTTP GET request. 271 func (s *TestServer) get(path string) *http.Response { 272 resp, err := s.HttpClient.Get(s.url(path)) 273 if err != nil { 274 s.t.Fatalf("err: %s", err) 275 } 276 if err := s.requireOK(resp); err != nil { 277 defer resp.Body.Close() 278 s.t.Fatal(err) 279 } 280 return resp 281 } 282 283 // encodePayload returns a new io.Reader wrapping the encoded contents 284 // of the payload, suitable for passing directly to a new request. 285 func (s *TestServer) encodePayload(payload interface{}) io.Reader { 286 var encoded bytes.Buffer 287 enc := json.NewEncoder(&encoded) 288 if err := enc.Encode(payload); err != nil { 289 s.t.Fatalf("err: %s", err) 290 } 291 return &encoded 292 }