github.com/huiliang/nomad@v0.2.1-0.20151124023127-7a8b664699ff/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  
    12  	"github.com/hashicorp/nomad/client/allocdir"
    13  	"github.com/hashicorp/nomad/client/config"
    14  	"github.com/hashicorp/nomad/nomad/structs"
    15  )
    16  
    17  var basicResources = &structs.Resources{
    18  	CPU:      250,
    19  	MemoryMB: 256,
    20  	Networks: []*structs.NetworkResource{
    21  		&structs.NetworkResource{
    22  			IP:            "0.0.0.0",
    23  			ReservedPorts: []structs.Port{{"main", 12345}},
    24  			DynamicPorts:  []structs.Port{{"HTTP", 43330}},
    25  		},
    26  	},
    27  }
    28  
    29  func init() {
    30  	rand.Seed(49875)
    31  }
    32  
    33  func testLogger() *log.Logger {
    34  	return log.New(os.Stderr, "", log.LstdFlags)
    35  }
    36  
    37  func testConfig() *config.Config {
    38  	conf := &config.Config{}
    39  	conf.StateDir = os.TempDir()
    40  	conf.AllocDir = os.TempDir()
    41  	return conf
    42  }
    43  
    44  func testDriverContext(task string) *DriverContext {
    45  	cfg := testConfig()
    46  	return NewDriverContext(task, cfg, cfg.Node, testLogger())
    47  }
    48  
    49  func testDriverExecContext(task *structs.Task, driverCtx *DriverContext) *ExecContext {
    50  	allocDir := allocdir.NewAllocDir(filepath.Join(driverCtx.config.AllocDir, structs.GenerateUUID()))
    51  	allocDir.Build([]*structs.Task{task})
    52  	ctx := NewExecContext(allocDir, fmt.Sprintf("alloc-id-%d", int(rand.Int31())))
    53  	return ctx
    54  }
    55  
    56  func TestDriver_TaskEnvironmentVariables(t *testing.T) {
    57  	ctx := &ExecContext{}
    58  	task := &structs.Task{
    59  		Env: map[string]string{
    60  			"HELLO": "world",
    61  			"lorem": "ipsum",
    62  		},
    63  		Resources: &structs.Resources{
    64  			CPU:      1000,
    65  			MemoryMB: 500,
    66  			Networks: []*structs.NetworkResource{
    67  				&structs.NetworkResource{
    68  					IP:            "1.2.3.4",
    69  					ReservedPorts: []structs.Port{{"one", 80}, {"two", 443}, {"three", 8080}, {"four", 12345}},
    70  					DynamicPorts:  []structs.Port{{"admin", 8081}, {"web", 8086}},
    71  				},
    72  			},
    73  		},
    74  		Meta: map[string]string{
    75  			"chocolate":  "cake",
    76  			"strawberry": "icecream",
    77  		},
    78  	}
    79  
    80  	env := TaskEnvironmentVariables(ctx, task)
    81  	exp := map[string]string{
    82  		"NOMAD_CPU_LIMIT":       "1000",
    83  		"NOMAD_MEMORY_LIMIT":    "500",
    84  		"NOMAD_IP":              "1.2.3.4",
    85  		"NOMAD_PORT_one":        "80",
    86  		"NOMAD_PORT_two":        "443",
    87  		"NOMAD_PORT_three":      "8080",
    88  		"NOMAD_PORT_four":       "12345",
    89  		"NOMAD_PORT_admin":      "8081",
    90  		"NOMAD_PORT_web":        "8086",
    91  		"NOMAD_META_CHOCOLATE":  "cake",
    92  		"NOMAD_META_STRAWBERRY": "icecream",
    93  		"HELLO":                 "world",
    94  		"lorem":                 "ipsum",
    95  	}
    96  
    97  	act := env.Map()
    98  	if !reflect.DeepEqual(act, exp) {
    99  		t.Fatalf("TaskEnvironmentVariables(%#v, %#v) returned %#v; want %#v", ctx, task, act, exp)
   100  	}
   101  }
   102  
   103  func TestMapMergeStrInt(t *testing.T) {
   104  	a := map[string]int{
   105  		"cakes":   5,
   106  		"cookies": 3,
   107  	}
   108  
   109  	b := map[string]int{
   110  		"cakes": 3,
   111  		"pies":  2,
   112  	}
   113  
   114  	c := mapMergeStrInt(a, b)
   115  
   116  	d := map[string]int{
   117  		"cakes":   3,
   118  		"cookies": 3,
   119  		"pies":    2,
   120  	}
   121  
   122  	if !reflect.DeepEqual(c, d) {
   123  		t.Errorf("\nExpected\n%+v\nGot\n%+v\n", d, c)
   124  	}
   125  }
   126  
   127  func TestMapMergeStrStr(t *testing.T) {
   128  	a := map[string]string{
   129  		"cake":   "chocolate",
   130  		"cookie": "caramel",
   131  	}
   132  
   133  	b := map[string]string{
   134  		"cake": "strawberry",
   135  		"pie":  "apple",
   136  	}
   137  
   138  	c := mapMergeStrStr(a, b)
   139  
   140  	d := map[string]string{
   141  		"cake":   "strawberry",
   142  		"cookie": "caramel",
   143  		"pie":    "apple",
   144  	}
   145  
   146  	if !reflect.DeepEqual(c, d) {
   147  		t.Errorf("\nExpected\n%+v\nGot\n%+v\n", d, c)
   148  	}
   149  }