Golang env.NewTaskEnvironment函数代码示例

// newTestHarness returns a harness starting a dev consul and vault server,
// building the appropriate config and creating a TaskTemplateManager
func newTestHarness(t *testing.T, templates []*structs.Template, consul, vault bool) *testHarness {
	harness := &testHarness{
		mockHooks: NewMockTaskHooks(),
		templates: templates,
		node:      mock.Node(),
		config:    &config.Config{},
	}

	// Build the task environment
	harness.taskEnv = env.NewTaskEnvironment(harness.node)

	// Make a tempdir
	d, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatalf("Failed to make tmpdir: %v", err)
	}
	harness.taskDir = d

	if consul {
		harness.consul = ctestutil.NewTestServer(t)
		harness.config.ConsulConfig = &sconfig.ConsulConfig{
			Addr: harness.consul.HTTPAddr,
		}
	}

	if vault {
		harness.vault = testutil.NewTestVault(t).Start()
		harness.config.VaultConfig = harness.vault.Config
		harness.vaultToken = harness.vault.RootToken
	}

	return harness
}

// GetTaskEnv converts the alloc dir, the node, task and alloc into a
// TaskEnvironment.
func GetTaskEnv(allocDir *allocdir.AllocDir, node *structs.Node,
	task *structs.Task, alloc *structs.Allocation) (*env.TaskEnvironment, error) {

	tg := alloc.Job.LookupTaskGroup(alloc.TaskGroup)
	env := env.NewTaskEnvironment(node).
		SetTaskMeta(task.Meta).
		SetTaskGroupMeta(tg.Meta).
		SetJobMeta(alloc.Job.Meta).
		SetEnvvars(task.Env).
		SetTaskName(task.Name)

	if allocDir != nil {
		env.SetAllocDir(allocDir.SharedDir)
		taskdir, ok := allocDir.TaskDirs[task.Name]
		if !ok {
			return nil, fmt.Errorf("failed to get task directory for task %q", task.Name)
		}

		env.SetTaskLocalDir(filepath.Join(taskdir, allocdir.TaskLocal))
	}

	if task.Resources != nil {
		env.SetMemLimit(task.Resources.MemoryMB).
			SetCpuLimit(task.Resources.CPU).
			SetNetworks(task.Resources.Networks)
	}

	if alloc != nil {
		env.SetAlloc(alloc)
	}

	return env.Build(), nil
}

func TestGetArtifact_File_RelativeDest(t *testing.T) {
	// Create the test server hosting the file to download
	ts := httptest.NewServer(http.FileServer(http.Dir(filepath.Dir("./test-fixtures/"))))
	defer ts.Close()

	// Create a temp directory to download into
	taskDir, err := ioutil.TempDir("", "nomad-test")
	if err != nil {
		t.Fatalf("failed to make temp directory: %v", err)
	}
	defer os.RemoveAll(taskDir)

	// Create the artifact
	file := "test.sh"
	relative := "foo/"
	artifact := &structs.TaskArtifact{
		GetterSource: fmt.Sprintf("%s/%s", ts.URL, file),
		GetterOptions: map[string]string{
			"checksum": "md5:bce963762aa2dbfed13caf492a45fb72",
		},
		RelativeDest: relative,
	}

	// Download the artifact
	taskEnv := env.NewTaskEnvironment(mock.Node())
	logger := log.New(os.Stderr, "", log.LstdFlags)
	if err := GetArtifact(taskEnv, artifact, taskDir, logger); err != nil {
		t.Fatalf("GetArtifact failed: %v", err)
	}

	// Verify artifact was downloaded to the correct path
	if _, err := os.Stat(filepath.Join(taskDir, relative, file)); err != nil {
		t.Fatalf("file not found: %s", err)
	}
}

func TestGetArtifact_InvalidChecksum(t *testing.T) {
	// Create the test server hosting the file to download
	ts := httptest.NewServer(http.FileServer(http.Dir(filepath.Dir("./test-fixtures/"))))
	defer ts.Close()

	// Create a temp directory to download into
	taskDir, err := ioutil.TempDir("", "nomad-test")
	if err != nil {
		t.Fatalf("failed to make temp directory: %v", err)
	}
	defer os.RemoveAll(taskDir)

	// Create the artifact with an incorrect checksum
	file := "test.sh"
	artifact := &structs.TaskArtifact{
		GetterSource: fmt.Sprintf("%s/%s", ts.URL, file),
		GetterOptions: map[string]string{
			"checksum": "md5:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
		},
	}

	// Download the artifact and expect an error
	taskEnv := env.NewTaskEnvironment(mock.Node())
	logger := log.New(os.Stderr, "", log.LstdFlags)
	if err := GetArtifact(taskEnv, artifact, taskDir, logger); err == nil {
		t.Fatalf("GetArtifact should have failed")
	}
}

// GetTaskEnv converts the alloc dir, the node and task configuration into a
// TaskEnvironment.
func GetTaskEnv(alloc *allocdir.AllocDir, node *structs.Node, task *structs.Task) (*env.TaskEnvironment, error) {
	env := env.NewTaskEnvironment(node).
		SetMeta(task.Meta).
		SetEnvvars(task.Env)

	if alloc != nil {
		env.SetAllocDir(alloc.SharedDir)
		taskdir, ok := alloc.TaskDirs[task.Name]
		if !ok {
			return nil, fmt.Errorf("failed to get task directory for task %q", task.Name)
		}

		env.SetTaskLocalDir(filepath.Join(taskdir, allocdir.TaskLocal))
	}

	if task.Resources != nil {
		env.SetMemLimit(task.Resources.MemoryMB)
		env.SetCpuLimit(task.Resources.CPU)

		if len(task.Resources.Networks) > 0 {
			network := task.Resources.Networks[0]
			env.SetTaskIp(network.IP)
			env.SetPorts(network.MapLabelToValues(nil))
		}
	}

	return env.Build(), nil
}

func testExecutorContext(t *testing.T) *ExecutorContext {
	taskEnv := env.NewTaskEnvironment(mock.Node())
	task, allocDir := mockAllocDir(t)
	ctx := &ExecutorContext{
		TaskEnv:  taskEnv,
		Task:     task,
		AllocDir: allocDir,
	}
	return ctx
}

func TestTaskTemplateManager_Invalid(t *testing.T) {
	hooks := NewMockTaskHooks()
	var tmpls []*structs.Template
	config := &config.Config{}
	taskDir := "foo"
	vaultToken := ""
	taskEnv := env.NewTaskEnvironment(mock.Node())

	_, err := NewTaskTemplateManager(nil, nil, nil, "", "", nil)
	if err == nil {
		t.Fatalf("Expected error")
	}

	_, err = NewTaskTemplateManager(nil, tmpls, config, vaultToken, taskDir, taskEnv)
	if err == nil || !strings.Contains(err.Error(), "task hook") {
		t.Fatalf("Expected invalid task hook error: %v", err)
	}

	_, err = NewTaskTemplateManager(hooks, tmpls, nil, vaultToken, taskDir, taskEnv)
	if err == nil || !strings.Contains(err.Error(), "config") {
		t.Fatalf("Expected invalid config error: %v", err)
	}

	_, err = NewTaskTemplateManager(hooks, tmpls, config, vaultToken, "", taskEnv)
	if err == nil || !strings.Contains(err.Error(), "task directory") {
		t.Fatalf("Expected invalid task dir error: %v", err)
	}

	_, err = NewTaskTemplateManager(hooks, tmpls, config, vaultToken, taskDir, nil)
	if err == nil || !strings.Contains(err.Error(), "task environment") {
		t.Fatalf("Expected invalid task environment error: %v", err)
	}

	tm, err := NewTaskTemplateManager(hooks, tmpls, config, vaultToken, taskDir, taskEnv)
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	} else if tm == nil {
		t.Fatalf("Bad %v", tm)
	}

	// Build a template with a bad signal
	tmpl := &structs.Template{
		DestPath:     "foo",
		EmbeddedTmpl: "hello, world",
		ChangeMode:   structs.TemplateChangeModeSignal,
		ChangeSignal: "foobarbaz",
	}

	tmpls = append(tmpls, tmpl)
	tm, err = NewTaskTemplateManager(hooks, tmpls, config, vaultToken, taskDir, taskEnv)
	if err == nil || !strings.Contains(err.Error(), "Failed to parse signal") {
		t.Fatalf("Expected signal parsing error: %v", err)
	}
}

func testExecutorContext(t *testing.T) *ExecutorContext {
	taskEnv := env.NewTaskEnvironment(mock.Node())
	taskName, allocDir := mockAllocDir(t)
	ctx := &ExecutorContext{
		TaskEnv:       taskEnv,
		TaskName:      taskName,
		AllocDir:      allocDir,
		TaskResources: constraint,
	}
	return ctx
}

func TestGetGetterUrl_Interprolation(t *testing.T) {
	// Create the artifact
	artifact := &structs.TaskArtifact{
		GetterSource: "${NOMAD_META_ARTIFACT}",
	}

	url := "foo.com"
	taskEnv := env.NewTaskEnvironment(mock.Node()).SetTaskMeta(map[string]string{"artifact": url})
	act, err := getGetterUrl(taskEnv, artifact)
	if err != nil {
		t.Fatalf("getGetterUrl() failed: %v", err)
	}

	if act != url {
		t.Fatalf("getGetterUrl() returned %q; want %q", act, url)
	}
}

func TestGetArtifact_Archive(t *testing.T) {
	// Create the test server hosting the file to download
	ts := httptest.NewServer(http.FileServer(http.Dir(filepath.Dir("./test-fixtures/"))))
	defer ts.Close()

	// Create a temp directory to download into and create some of the same
	// files that exist in the artifact to ensure they are overridden
	taskDir, err := ioutil.TempDir("", "nomad-test")
	if err != nil {
		t.Fatalf("failed to make temp directory: %v", err)
	}
	defer os.RemoveAll(taskDir)

	create := map[string]string{
		"exist/my.config": "to be replaced",
		"untouched":       "existing top-level",
	}
	createContents(taskDir, create, t)

	file := "archive.tar.gz"
	artifact := &structs.TaskArtifact{
		GetterSource: fmt.Sprintf("%s/%s", ts.URL, file),
		GetterOptions: map[string]string{
			"checksum": "sha1:20bab73c72c56490856f913cf594bad9a4d730f6",
		},
	}

	taskEnv := env.NewTaskEnvironment(mock.Node())
	logger := log.New(os.Stderr, "", log.LstdFlags)
	if err := GetArtifact(taskEnv, artifact, taskDir, logger); err != nil {
		t.Fatalf("GetArtifact failed: %v", err)
	}

	// Verify the unarchiving overrode files properly.
	expected := map[string]string{
		"untouched":       "existing top-level",
		"exist/my.config": "hello world\n",
		"new/my.config":   "hello world\n",
		"test.sh":         "sleep 1\n",
	}
	checkContents(taskDir, expected, t)
}

func testExecutorContextWithChroot(t *testing.T) *ExecutorContext {
	taskEnv := env.NewTaskEnvironment(mock.Node())
	task, allocDir := mockAllocDir(t)
	ctx := &ExecutorContext{
		TaskEnv:  taskEnv,
		Task:     task,
		AllocDir: allocDir,
		ChrootEnv: map[string]string{
			"/etc/ld.so.cache":  "/etc/ld.so.cache",
			"/etc/ld.so.conf":   "/etc/ld.so.conf",
			"/etc/ld.so.conf.d": "/etc/ld.so.conf.d",
			"/lib":              "/lib",
			"/lib64":            "/lib64",
			"/usr/lib":          "/usr/lib",
			"/bin/ls":           "/bin/ls",
			"/foobar":           "/does/not/exist",
		},
	}
	return ctx
}

func testExecutorContext() *ExecutorContext {
	taskEnv := env.NewTaskEnvironment(mock.Node())
	return &ExecutorContext{taskEnv: taskEnv}
}