github.com/ranjib/nomad@v0.1.1-0.20160225204057-97751b02f70b/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 conf.MaxKillTimeout = 10 * time.Second 50 return conf 51 } 52 53 func testDriverContexts(task *structs.Task) (*DriverContext, *ExecContext) { 54 cfg := testConfig() 55 allocDir := allocdir.NewAllocDir(filepath.Join(cfg.AllocDir, structs.GenerateUUID())) 56 allocDir.Build([]*structs.Task{task}) 57 execCtx := NewExecContext(allocDir, fmt.Sprintf("alloc-id-%d", int(rand.Int31()))) 58 59 taskEnv, err := GetTaskEnv(allocDir, cfg.Node, task) 60 if err != nil { 61 return nil, nil 62 } 63 64 driverCtx := NewDriverContext(task.Name, cfg, cfg.Node, testLogger(), taskEnv) 65 return driverCtx, execCtx 66 } 67 68 func TestDriver_KillTimeout(t *testing.T) { 69 expected := 1 * time.Second 70 task := &structs.Task{Name: "foo", KillTimeout: expected} 71 ctx, _ := testDriverContexts(task) 72 ctx.config.MaxKillTimeout = 10 * time.Second 73 74 if actual := ctx.KillTimeout(task); expected != actual { 75 t.Fatalf("KillTimeout(%v) returned %v; want %v", task, actual, expected) 76 } 77 78 expected = 10 * time.Second 79 task = &structs.Task{KillTimeout: 11 * time.Second} 80 81 if actual := ctx.KillTimeout(task); expected != actual { 82 t.Fatalf("KillTimeout(%v) returned %v; want %v", task, actual, expected) 83 } 84 } 85 86 func TestDriver_GetTaskEnv(t *testing.T) { 87 t.Parallel() 88 task := &structs.Task{ 89 Env: map[string]string{ 90 "HELLO": "world", 91 "lorem": "ipsum", 92 }, 93 Resources: &structs.Resources{ 94 CPU: 1000, 95 MemoryMB: 500, 96 Networks: []*structs.NetworkResource{ 97 &structs.NetworkResource{ 98 IP: "1.2.3.4", 99 ReservedPorts: []structs.Port{{"one", 80}, {"two", 443}, {"three", 8080}, {"four", 12345}}, 100 DynamicPorts: []structs.Port{{"admin", 8081}, {"web", 8086}}, 101 }, 102 }, 103 }, 104 Meta: map[string]string{ 105 "chocolate": "cake", 106 "strawberry": "icecream", 107 }, 108 } 109 110 env, err := GetTaskEnv(nil, nil, task) 111 if err != nil { 112 t.Fatalf("GetTaskEnv() failed: %v", err) 113 } 114 exp := map[string]string{ 115 "NOMAD_CPU_LIMIT": "1000", 116 "NOMAD_MEMORY_LIMIT": "500", 117 "NOMAD_ADDR_one": "1.2.3.4:80", 118 "NOMAD_ADDR_two": "1.2.3.4:443", 119 "NOMAD_ADDR_three": "1.2.3.4:8080", 120 "NOMAD_ADDR_four": "1.2.3.4:12345", 121 "NOMAD_ADDR_admin": "1.2.3.4:8081", 122 "NOMAD_ADDR_web": "1.2.3.4: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 }