github.com/dkerwin/nomad@v0.3.3-0.20160525181927-74554135514b/client/driver/driver_test.go (about) 1 package driver 2 3 import ( 4 "io" 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/mock" 17 "github.com/hashicorp/nomad/nomad/structs" 18 ) 19 20 var basicResources = &structs.Resources{ 21 CPU: 250, 22 MemoryMB: 256, 23 Networks: []*structs.NetworkResource{ 24 &structs.NetworkResource{ 25 IP: "0.0.0.0", 26 ReservedPorts: []structs.Port{{"main", 12345}}, 27 DynamicPorts: []structs.Port{{"HTTP", 43330}}, 28 }, 29 }, 30 } 31 32 func init() { 33 rand.Seed(49875) 34 } 35 36 func TestMain(m *testing.M) { 37 if !testtask.Run() { 38 os.Exit(m.Run()) 39 } 40 } 41 42 // copyFile moves an existing file to the destination 43 func copyFile(src, dst string, t *testing.T) { 44 in, err := os.Open(src) 45 if err != nil { 46 t.Fatalf("copying %v -> %v failed: %v", src, dst, err) 47 } 48 defer in.Close() 49 out, err := os.Create(dst) 50 if err != nil { 51 t.Fatalf("copying %v -> %v failed: %v", src, dst, err) 52 } 53 defer func() { 54 if err := out.Close(); err != nil { 55 t.Fatalf("copying %v -> %v failed: %v", src, dst, err) 56 } 57 }() 58 if _, err = io.Copy(out, in); err != nil { 59 t.Fatalf("copying %v -> %v failed: %v", src, dst, err) 60 } 61 if err := out.Sync(); err != nil { 62 t.Fatalf("copying %v -> %v failed: %v", src, dst, err) 63 } 64 return 65 } 66 67 func testLogger() *log.Logger { 68 return log.New(os.Stderr, "", log.LstdFlags) 69 } 70 71 func testConfig() *config.Config { 72 conf := &config.Config{} 73 conf.StateDir = os.TempDir() 74 conf.AllocDir = os.TempDir() 75 conf.MaxKillTimeout = 10 * time.Second 76 return conf 77 } 78 79 func testDriverContexts(task *structs.Task) (*DriverContext, *ExecContext) { 80 cfg := testConfig() 81 allocDir := allocdir.NewAllocDir(filepath.Join(cfg.AllocDir, structs.GenerateUUID())) 82 allocDir.Build([]*structs.Task{task}) 83 alloc := mock.Alloc() 84 execCtx := NewExecContext(allocDir, alloc.ID) 85 86 taskEnv, err := GetTaskEnv(allocDir, cfg.Node, task, alloc) 87 if err != nil { 88 return nil, nil 89 } 90 91 driverCtx := NewDriverContext(task.Name, cfg, cfg.Node, testLogger(), taskEnv) 92 return driverCtx, execCtx 93 } 94 95 func TestDriver_GetTaskEnv(t *testing.T) { 96 t.Parallel() 97 task := &structs.Task{ 98 Name: "Foo", 99 Env: map[string]string{ 100 "HELLO": "world", 101 "lorem": "ipsum", 102 }, 103 Resources: &structs.Resources{ 104 CPU: 1000, 105 MemoryMB: 500, 106 Networks: []*structs.NetworkResource{ 107 &structs.NetworkResource{ 108 IP: "1.2.3.4", 109 ReservedPorts: []structs.Port{{"one", 80}, {"two", 443}}, 110 DynamicPorts: []structs.Port{{"admin", 8081}, {"web", 8086}}, 111 }, 112 }, 113 }, 114 Meta: map[string]string{ 115 "chocolate": "cake", 116 "strawberry": "icecream", 117 }, 118 } 119 120 alloc := mock.Alloc() 121 alloc.Name = "Bar" 122 env, err := GetTaskEnv(nil, nil, task, alloc) 123 if err != nil { 124 t.Fatalf("GetTaskEnv() failed: %v", err) 125 } 126 exp := map[string]string{ 127 "NOMAD_CPU_LIMIT": "1000", 128 "NOMAD_MEMORY_LIMIT": "500", 129 "NOMAD_ADDR_one": "1.2.3.4:80", 130 "NOMAD_IP_one": "1.2.3.4", 131 "NOMAD_PORT_one": "80", 132 "NOMAD_ADDR_two": "1.2.3.4:443", 133 "NOMAD_IP_two": "1.2.3.4", 134 "NOMAD_PORT_two": "443", 135 "NOMAD_ADDR_admin": "1.2.3.4:8081", 136 "NOMAD_IP_admin": "1.2.3.4", 137 "NOMAD_PORT_admin": "8081", 138 "NOMAD_ADDR_web": "1.2.3.4:8086", 139 "NOMAD_IP_web": "1.2.3.4", 140 "NOMAD_PORT_web": "8086", 141 "NOMAD_META_CHOCOLATE": "cake", 142 "NOMAD_META_STRAWBERRY": "icecream", 143 "NOMAD_META_ELB_CHECK_INTERVAL": "30s", 144 "NOMAD_META_ELB_CHECK_TYPE": "http", 145 "NOMAD_META_ELB_CHECK_MIN": "3", 146 "NOMAD_META_OWNER": "armon", 147 "HELLO": "world", 148 "lorem": "ipsum", 149 "NOMAD_ALLOC_ID": alloc.ID, 150 "NOMAD_ALLOC_NAME": alloc.Name, 151 "NOMAD_TASK_NAME": task.Name, 152 } 153 154 act := env.EnvMap() 155 if !reflect.DeepEqual(act, exp) { 156 t.Fatalf("GetTaskEnv() returned %#v; want %#v", act, exp) 157 } 158 } 159 160 func TestMapMergeStrInt(t *testing.T) { 161 t.Parallel() 162 a := map[string]int{ 163 "cakes": 5, 164 "cookies": 3, 165 } 166 167 b := map[string]int{ 168 "cakes": 3, 169 "pies": 2, 170 } 171 172 c := mapMergeStrInt(a, b) 173 174 d := map[string]int{ 175 "cakes": 3, 176 "cookies": 3, 177 "pies": 2, 178 } 179 180 if !reflect.DeepEqual(c, d) { 181 t.Errorf("\nExpected\n%+v\nGot\n%+v\n", d, c) 182 } 183 } 184 185 func TestMapMergeStrStr(t *testing.T) { 186 t.Parallel() 187 a := map[string]string{ 188 "cake": "chocolate", 189 "cookie": "caramel", 190 } 191 192 b := map[string]string{ 193 "cake": "strawberry", 194 "pie": "apple", 195 } 196 197 c := mapMergeStrStr(a, b) 198 199 d := map[string]string{ 200 "cake": "strawberry", 201 "cookie": "caramel", 202 "pie": "apple", 203 } 204 205 if !reflect.DeepEqual(c, d) { 206 t.Errorf("\nExpected\n%+v\nGot\n%+v\n", d, c) 207 } 208 }