github.com/blixtra/nomad@v0.7.2-0.20171221000451-da9a1d7bb050/command/agent/testagent.go (about) 1 package agent 2 3 import ( 4 "fmt" 5 "io" 6 "io/ioutil" 7 "math/rand" 8 "net/http" 9 "net/http/httptest" 10 "os" 11 "path/filepath" 12 "runtime" 13 "strings" 14 "time" 15 16 "github.com/mitchellh/go-testing-interface" 17 18 metrics "github.com/armon/go-metrics" 19 "github.com/hashicorp/consul/lib/freeport" 20 "github.com/hashicorp/nomad/api" 21 "github.com/hashicorp/nomad/client/fingerprint" 22 "github.com/hashicorp/nomad/nomad" 23 "github.com/hashicorp/nomad/nomad/mock" 24 "github.com/hashicorp/nomad/nomad/structs" 25 sconfig "github.com/hashicorp/nomad/nomad/structs/config" 26 "github.com/hashicorp/nomad/testutil" 27 ) 28 29 func init() { 30 rand.Seed(time.Now().UnixNano()) // seed random number generator 31 } 32 33 // TempDir defines the base dir for temporary directories. 34 var TempDir = os.TempDir() 35 36 // TestAgent encapsulates an Agent with a default configuration and startup 37 // procedure suitable for testing. It manages a temporary data directory which 38 // is removed after shutdown. 39 type TestAgent struct { 40 // T is the testing object 41 T testing.T 42 43 // Name is an optional name of the agent. 44 Name string 45 46 // ConfigCallback is an optional callback that allows modification of the 47 // configuration before the agent is started. 48 ConfigCallback func(*Config) 49 50 // Config is the agent configuration. If Config is nil then 51 // TestConfig() is used. If Config.DataDir is set then it is 52 // the callers responsibility to clean up the data directory. 53 // Otherwise, a temporary data directory is created and removed 54 // when Shutdown() is called. 55 Config *Config 56 57 // LogOutput is the sink for the logs. If nil, logs are written 58 // to os.Stderr. 59 LogOutput io.Writer 60 61 // DataDir is the data directory which is used when Config.DataDir 62 // is not set. It is created automatically and removed when 63 // Shutdown() is called. 64 DataDir string 65 66 // Key is the optional encryption key for the keyring. 67 Key string 68 69 // Server is a reference to the started HTTP endpoint. 70 // It is valid after Start(). 71 Server *HTTPServer 72 73 // Agent is the embedded Nomad agent. 74 // It is valid after Start(). 75 *Agent 76 77 // RootToken is auto-bootstrapped if ACLs are enabled 78 RootToken *structs.ACLToken 79 } 80 81 // NewTestAgent returns a started agent with the given name and 82 // configuration. The caller should call Shutdown() to stop the agent and 83 // remove temporary directories. 84 func NewTestAgent(t testing.T, name string, configCallback func(*Config)) *TestAgent { 85 a := &TestAgent{ 86 T: t, 87 Name: name, 88 ConfigCallback: configCallback, 89 } 90 91 a.Start() 92 return a 93 } 94 95 // Start starts a test agent. 96 func (a *TestAgent) Start() *TestAgent { 97 if a.Agent != nil { 98 a.T.Fatalf("TestAgent already started") 99 } 100 if a.Config == nil { 101 a.Config = a.config() 102 } 103 if a.Config.DataDir == "" { 104 name := "agent" 105 if a.Name != "" { 106 name = a.Name + "-agent" 107 } 108 name = strings.Replace(name, "/", "_", -1) 109 d, err := ioutil.TempDir(TempDir, name) 110 if err != nil { 111 a.T.Fatalf("Error creating data dir %s: %s", filepath.Join(TempDir, name), err) 112 } 113 a.DataDir = d 114 a.Config.DataDir = d 115 a.Config.NomadConfig.DataDir = d 116 } 117 118 for i := 10; i >= 0; i-- { 119 a.pickRandomPorts(a.Config) 120 if a.Config.NodeName == "" { 121 a.Config.NodeName = fmt.Sprintf("Node %d", a.Config.Ports.RPC) 122 } 123 124 // write the keyring 125 if a.Key != "" { 126 writeKey := func(key, filename string) { 127 path := filepath.Join(a.Config.DataDir, filename) 128 if err := initKeyring(path, key); err != nil { 129 a.T.Fatalf("Error creating keyring %s: %s", path, err) 130 } 131 } 132 writeKey(a.Key, serfKeyring) 133 } 134 135 // we need the err var in the next exit condition 136 if agent, err := a.start(); err == nil { 137 a.Agent = agent 138 break 139 } else if i == 0 { 140 fmt.Println(a.Name, "Error starting agent:", err) 141 runtime.Goexit() 142 } else { 143 if agent != nil { 144 agent.Shutdown() 145 } 146 wait := time.Duration(rand.Int31n(2000)) * time.Millisecond 147 fmt.Println(a.Name, "retrying in", wait) 148 time.Sleep(wait) 149 } 150 151 // Clean out the data dir if we are responsible for it before we 152 // try again, since the old ports may have gotten written to 153 // the data dir, such as in the Raft configuration. 154 if a.DataDir != "" { 155 if err := os.RemoveAll(a.DataDir); err != nil { 156 fmt.Println(a.Name, "Error resetting data dir:", err) 157 runtime.Goexit() 158 } 159 } 160 } 161 162 if a.Config.NomadConfig.Bootstrap && a.Config.Server.Enabled { 163 testutil.WaitForResult(func() (bool, error) { 164 args := &structs.GenericRequest{} 165 var leader string 166 err := a.RPC("Status.Leader", args, &leader) 167 return leader != "", err 168 }, func(err error) { 169 a.T.Fatalf("failed to find leader: %v", err) 170 }) 171 } else { 172 testutil.WaitForResult(func() (bool, error) { 173 req, _ := http.NewRequest("GET", "/v1/agent/self", nil) 174 resp := httptest.NewRecorder() 175 _, err := a.Server.AgentSelfRequest(resp, req) 176 return err == nil && resp.Code == 200, err 177 }, func(err error) { 178 a.T.Fatalf("failed OK response: %v", err) 179 }) 180 } 181 182 // Check if ACLs enabled. Use special value of PolicyTTL 0s 183 // to do a bypass of this step. This is so we can test bootstrap 184 // without having to pass down a special flag. 185 if a.Config.ACL.Enabled && a.Config.Server.Enabled && a.Config.ACL.PolicyTTL != 0 { 186 a.RootToken = mock.ACLManagementToken() 187 state := a.Agent.server.State() 188 if err := state.BootstrapACLTokens(1, 0, a.RootToken); err != nil { 189 a.T.Fatalf("token bootstrap failed: %v", err) 190 } 191 } 192 return a 193 } 194 195 func (a *TestAgent) start() (*Agent, error) { 196 if a.LogOutput == nil { 197 a.LogOutput = os.Stderr 198 } 199 200 inm := metrics.NewInmemSink(10*time.Second, time.Minute) 201 metrics.NewGlobal(metrics.DefaultConfig("service-name"), inm) 202 203 if inm == nil { 204 return nil, fmt.Errorf("unable to set up in memory metrics needed for agent initialization") 205 } 206 207 agent, err := NewAgent(a.Config, a.LogOutput, inm) 208 if err != nil { 209 return nil, err 210 } 211 212 // Setup the HTTP server 213 http, err := NewHTTPServer(agent, a.Config) 214 if err != nil { 215 return agent, err 216 } 217 218 a.Server = http 219 return agent, nil 220 } 221 222 // Shutdown stops the agent and removes the data directory if it is 223 // managed by the test agent. 224 func (a *TestAgent) Shutdown() error { 225 defer func() { 226 if a.DataDir != "" { 227 os.RemoveAll(a.DataDir) 228 } 229 }() 230 231 // shutdown agent before endpoints 232 a.Server.Shutdown() 233 return a.Agent.Shutdown() 234 } 235 236 func (a *TestAgent) HTTPAddr() string { 237 if a.Server == nil { 238 return "" 239 } 240 return "http://" + a.Server.Addr 241 } 242 243 func (a *TestAgent) Client() *api.Client { 244 conf := api.DefaultConfig() 245 conf.Address = a.HTTPAddr() 246 c, err := api.NewClient(conf) 247 if err != nil { 248 a.T.Fatalf("Error creating Nomad API client: %s", err) 249 } 250 return c 251 } 252 253 // pickRandomPorts selects random ports from fixed size random blocks of 254 // ports. This does not eliminate the chance for port conflict but 255 // reduces it significanltly with little overhead. Furthermore, asking 256 // the kernel for a random port by binding to port 0 prolongs the test 257 // execution (in our case +20sec) while also not fully eliminating the 258 // chance of port conflicts for concurrently executed test binaries. 259 // Instead of relying on one set of ports to be sufficient we retry 260 // starting the agent with different ports on port conflict. 261 func (a *TestAgent) pickRandomPorts(c *Config) { 262 ports := freeport.GetT(a.T, 3) 263 c.Ports.HTTP = ports[0] 264 c.Ports.RPC = ports[1] 265 c.Ports.Serf = ports[2] 266 267 if err := c.normalizeAddrs(); err != nil { 268 a.T.Fatalf("error normalizing config: %v", err) 269 } 270 } 271 272 // TestConfig returns a unique default configuration for testing an 273 // agent. 274 func (a *TestAgent) config() *Config { 275 conf := DevConfig() 276 277 // Customize the server configuration 278 config := nomad.DefaultConfig() 279 conf.NomadConfig = config 280 281 // Set the name 282 conf.NodeName = a.Name 283 284 // Bind and set ports 285 conf.BindAddr = "127.0.0.1" 286 287 conf.Consul = sconfig.DefaultConsulConfig() 288 conf.Vault.Enabled = new(bool) 289 290 // Tighten the Serf timing 291 config.SerfConfig.MemberlistConfig.SuspicionMult = 2 292 config.SerfConfig.MemberlistConfig.RetransmitMult = 2 293 config.SerfConfig.MemberlistConfig.ProbeTimeout = 50 * time.Millisecond 294 config.SerfConfig.MemberlistConfig.ProbeInterval = 100 * time.Millisecond 295 config.SerfConfig.MemberlistConfig.GossipInterval = 100 * time.Millisecond 296 297 // Tighten the Raft timing 298 config.RaftConfig.LeaderLeaseTimeout = 20 * time.Millisecond 299 config.RaftConfig.HeartbeatTimeout = 40 * time.Millisecond 300 config.RaftConfig.ElectionTimeout = 40 * time.Millisecond 301 config.RaftConfig.StartAsLeader = true 302 config.RaftTimeout = 500 * time.Millisecond 303 304 // Bootstrap ourselves 305 config.Bootstrap = true 306 config.BootstrapExpect = 1 307 308 // Tighten the fingerprinter timeouts 309 if conf.Client.Options == nil { 310 conf.Client.Options = make(map[string]string) 311 } 312 conf.Client.Options[fingerprint.TightenNetworkTimeoutsConfig] = "true" 313 314 if a.ConfigCallback != nil { 315 a.ConfigCallback(conf) 316 } 317 318 return conf 319 }