124 lines
3.0 KiB
Go
124 lines
3.0 KiB
Go
package command
|
|
|
|
import (
|
|
"os"
|
|
"testing"
|
|
|
|
"github.com/hashicorp/nomad/api"
|
|
"github.com/hashicorp/nomad/command/agent"
|
|
"github.com/hashicorp/nomad/helper"
|
|
"github.com/hashicorp/nomad/testutil"
|
|
)
|
|
|
|
func testServer(t *testing.T, runClient bool, cb func(*agent.Config)) (*agent.TestAgent, *api.Client, string) {
|
|
// Make a new test server
|
|
a := agent.NewTestAgent(t, t.Name(), func(config *agent.Config) {
|
|
config.Client.Enabled = runClient
|
|
|
|
if cb != nil {
|
|
cb(config)
|
|
}
|
|
})
|
|
t.Cleanup(func() { a.Shutdown() })
|
|
|
|
c := a.Client()
|
|
return a, c, a.HTTPAddr()
|
|
}
|
|
|
|
// testClient starts a new test client, blocks until it joins, and performs
|
|
// cleanup after the test is complete.
|
|
func testClient(t *testing.T, name string, cb func(*agent.Config)) (*agent.TestAgent, *api.Client, string) {
|
|
t.Logf("[TEST] Starting client agent %s", name)
|
|
a := agent.NewTestAgent(t, name, func(config *agent.Config) {
|
|
if cb != nil {
|
|
cb(config)
|
|
}
|
|
})
|
|
t.Cleanup(func() { a.Shutdown() })
|
|
|
|
c := a.Client()
|
|
t.Logf("[TEST] Waiting for client %s to join server(s) %s", name, a.GetConfig().Client.Servers)
|
|
testutil.WaitForClient(t, a.Agent.RPC, a.Agent.Client().NodeID(), a.Agent.Client().Region())
|
|
|
|
return a, c, a.HTTPAddr()
|
|
}
|
|
|
|
func testJob(jobID string) *api.Job {
|
|
task := api.NewTask("task1", "mock_driver").
|
|
SetConfig("kill_after", "1s").
|
|
SetConfig("run_for", "5s").
|
|
SetConfig("exit_code", 0).
|
|
Require(&api.Resources{
|
|
MemoryMB: helper.IntToPtr(256),
|
|
CPU: helper.IntToPtr(100),
|
|
}).
|
|
SetLogConfig(&api.LogConfig{
|
|
MaxFiles: helper.IntToPtr(1),
|
|
MaxFileSizeMB: helper.IntToPtr(2),
|
|
})
|
|
|
|
group := api.NewTaskGroup("group1", 1).
|
|
AddTask(task).
|
|
RequireDisk(&api.EphemeralDisk{
|
|
SizeMB: helper.IntToPtr(20),
|
|
})
|
|
|
|
job := api.NewBatchJob(jobID, jobID, "global", 1).
|
|
AddDatacenter("dc1").
|
|
AddTaskGroup(group)
|
|
|
|
return job
|
|
}
|
|
|
|
func testMultiRegionJob(jobID, region, datacenter string) *api.Job {
|
|
task := api.NewTask("task1", "mock_driver").
|
|
SetConfig("kill_after", "10s").
|
|
SetConfig("run_for", "15s").
|
|
SetConfig("exit_code", 0).
|
|
Require(&api.Resources{
|
|
MemoryMB: helper.IntToPtr(256),
|
|
CPU: helper.IntToPtr(100),
|
|
}).
|
|
SetLogConfig(&api.LogConfig{
|
|
MaxFiles: helper.IntToPtr(1),
|
|
MaxFileSizeMB: helper.IntToPtr(2),
|
|
})
|
|
|
|
group := api.NewTaskGroup("group1", 1).
|
|
AddTask(task).
|
|
RequireDisk(&api.EphemeralDisk{
|
|
SizeMB: helper.IntToPtr(20),
|
|
})
|
|
|
|
job := api.NewServiceJob(jobID, jobID, region, 1).AddDatacenter(datacenter).AddTaskGroup(group)
|
|
job.Region = nil
|
|
job.Multiregion = &api.Multiregion{
|
|
Regions: []*api.MultiregionRegion{
|
|
{
|
|
Name: "east",
|
|
Datacenters: []string{"east-1"},
|
|
},
|
|
{
|
|
Name: "west",
|
|
Datacenters: []string{"west-1"},
|
|
},
|
|
},
|
|
}
|
|
|
|
return job
|
|
}
|
|
|
|
// setEnv wraps os.Setenv(key, value) and restores the environment variable to initial value in test cleanup
|
|
func setEnv(t *testing.T, key, value string) {
|
|
initial, ok := os.LookupEnv(key)
|
|
os.Setenv(key, value)
|
|
|
|
t.Cleanup(func() {
|
|
if ok {
|
|
os.Setenv(key, initial)
|
|
} else {
|
|
os.Unsetenv(key)
|
|
}
|
|
})
|
|
}
|