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  }