github.com/ryanslade/nomad@v0.2.4-0.20160128061903-fc95782f2089/client/driver/driver_test.go (about) 1 package driver 2 3 import ( 4 "fmt" 5 "log" 6 "math/rand" 7 "os" 8 "path/filepath" 9 "reflect" 10 "testing" 11 "time" 12 13 "github.com/hashicorp/nomad/client/allocdir" 14 "github.com/hashicorp/nomad/client/config" 15 "github.com/hashicorp/nomad/helper/testtask" 16 "github.com/hashicorp/nomad/nomad/structs" 17 ) 18 19 var basicResources = &structs.Resources{ 20 CPU: 250, 21 MemoryMB: 256, 22 Networks: []*structs.NetworkResource{ 23 &structs.NetworkResource{ 24 IP: "0.0.0.0", 25 ReservedPorts: []structs.Port{{"main", 12345}}, 26 DynamicPorts: []structs.Port{{"HTTP", 43330}}, 27 }, 28 }, 29 } 30 31 func init() { 32 rand.Seed(49875) 33 } 34 35 func TestMain(m *testing.M) { 36 if !testtask.Run() { 37 os.Exit(m.Run()) 38 } 39 } 40 41 func testLogger() *log.Logger { 42 return log.New(os.Stderr, "", log.LstdFlags) 43 } 44 45 func testConfig() *config.Config { 46 conf := &config.Config{} 47 conf.StateDir = os.TempDir() 48 conf.AllocDir = os.TempDir() 49 return conf 50 } 51 52 func testDriverContexts(task *structs.Task) (*DriverContext, *ExecContext) { 53 cfg := testConfig() 54 allocDir := allocdir.NewAllocDir(filepath.Join(cfg.AllocDir, structs.GenerateUUID())) 55 allocDir.Build([]*structs.Task{task}) 56 execCtx := NewExecContext(allocDir, fmt.Sprintf("alloc-id-%d", int(rand.Int31()))) 57 58 taskEnv, err := GetTaskEnv(allocDir, cfg.Node, task) 59 if err != nil { 60 return nil, nil 61 } 62 63 driverCtx := NewDriverContext(task.Name, cfg, cfg.Node, testLogger(), taskEnv) 64 return driverCtx, execCtx 65 } 66 67 func TestDriver_KillTimeout(t *testing.T) { 68 expected := 1 * time.Second 69 task := &structs.Task{Name: "foo", KillTimeout: expected} 70 ctx, _ := testDriverContexts(task) 71 ctx.config.MaxKillTimeout = 10 * time.Second 72 73 if actual := ctx.KillTimeout(task); expected != actual { 74 t.Fatalf("KillTimeout(%v) returned %v; want %v", task, actual, expected) 75 } 76 77 expected = 10 * time.Second 78 task = &structs.Task{KillTimeout: 11 * time.Second} 79 80 if actual := ctx.KillTimeout(task); expected != actual { 81 t.Fatalf("KillTimeout(%v) returned %v; want %v", task, actual, expected) 82 } 83 } 84 85 func TestDriver_GetTaskEnv(t *testing.T) { 86 t.Parallel() 87 task := &structs.Task{ 88 Env: map[string]string{ 89 "HELLO": "world", 90 "lorem": "ipsum", 91 }, 92 Resources: &structs.Resources{ 93 CPU: 1000, 94 MemoryMB: 500, 95 Networks: []*structs.NetworkResource{ 96 &structs.NetworkResource{ 97 IP: "1.2.3.4", 98 ReservedPorts: []structs.Port{{"one", 80}, {"two", 443}, {"three", 8080}, {"four", 12345}}, 99 DynamicPorts: []structs.Port{{"admin", 8081}, {"web", 8086}}, 100 }, 101 }, 102 }, 103 Meta: map[string]string{ 104 "chocolate": "cake", 105 "strawberry": "icecream", 106 }, 107 } 108 109 env, err := GetTaskEnv(nil, nil, task) 110 if err != nil { 111 t.Fatalf("GetTaskEnv() failed: %v", err) 112 } 113 exp := map[string]string{ 114 "NOMAD_CPU_LIMIT": "1000", 115 "NOMAD_MEMORY_LIMIT": "500", 116 "NOMAD_IP": "1.2.3.4", 117 "NOMAD_PORT_one": "80", 118 "NOMAD_PORT_two": "443", 119 "NOMAD_PORT_three": "8080", 120 "NOMAD_PORT_four": "12345", 121 "NOMAD_PORT_admin": "8081", 122 "NOMAD_PORT_web": "8086", 123 "NOMAD_META_CHOCOLATE": "cake", 124 "NOMAD_META_STRAWBERRY": "icecream", 125 "HELLO": "world", 126 "lorem": "ipsum", 127 } 128 129 act := env.EnvMap() 130 if !reflect.DeepEqual(act, exp) { 131 t.Fatalf("GetTaskEnv() returned %#v; want %#v", act, exp) 132 } 133 } 134 135 func TestMapMergeStrInt(t *testing.T) { 136 t.Parallel() 137 a := map[string]int{ 138 "cakes": 5, 139 "cookies": 3, 140 } 141 142 b := map[string]int{ 143 "cakes": 3, 144 "pies": 2, 145 } 146 147 c := mapMergeStrInt(a, b) 148 149 d := map[string]int{ 150 "cakes": 3, 151 "cookies": 3, 152 "pies": 2, 153 } 154 155 if !reflect.DeepEqual(c, d) { 156 t.Errorf("\nExpected\n%+v\nGot\n%+v\n", d, c) 157 } 158 } 159 160 func TestMapMergeStrStr(t *testing.T) { 161 t.Parallel() 162 a := map[string]string{ 163 "cake": "chocolate", 164 "cookie": "caramel", 165 } 166 167 b := map[string]string{ 168 "cake": "strawberry", 169 "pie": "apple", 170 } 171 172 c := mapMergeStrStr(a, b) 173 174 d := map[string]string{ 175 "cake": "strawberry", 176 "cookie": "caramel", 177 "pie": "apple", 178 } 179 180 if !reflect.DeepEqual(c, d) { 181 t.Errorf("\nExpected\n%+v\nGot\n%+v\n", d, c) 182 } 183 }