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  }