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  }