<commit_before>package drivers

import (
	"context"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"runtime"
	"time"

	hclog "github.com/hashicorp/go-hclog"
	plugin "github.com/hashicorp/go-plugin"
	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/driver/env"
	"github.com/hashicorp/nomad/client/logmon"
	cstructs "github.com/hashicorp/nomad/client/structs"
	"github.com/hashicorp/nomad/helper/testlog"
	"github.com/hashicorp/nomad/helper/uuid"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/plugins/base"
	"github.com/hashicorp/nomad/plugins/drivers/utils"
	"github.com/hashicorp/nomad/plugins/shared/hclspec"
	"github.com/mitchellh/go-testing-interface"
	"github.com/stretchr/testify/require"
)

type DriverHarness struct {
	DriverPlugin
	client *plugin.GRPCClient
	server *plugin.GRPCServer
	t      testing.T
	lm     logmon.LogMon
	logger hclog.Logger
	impl   DriverPlugin
}

func (d *DriverHarness) Impl() DriverPlugin {
	return d.impl
}
func NewDriverHarness(t testing.T, d DriverPlugin) *DriverHarness {
	logger := testlog.HCLogger(t).Named("driver_harness")
	client, server := plugin.TestPluginGRPCConn(t,
		map[string]plugin.Plugin{
			base.PluginTypeDriver: &PluginDriver{
				impl:   d,
				logger: logger.Named("driver_plugin"),
			},
			base.PluginTypeBase: &base.PluginBase{Impl: d},
			"logmon":            logmon.NewPlugin(logmon.NewLogMon(logger.Named("logmon"))),
		},
	)

	raw, err := client.Dispense(base.PluginTypeDriver)
	if err != nil {
		t.Fatalf("err dispensing plugin: %v", err)
	}

	dClient := raw.(DriverPlugin)
	h := &DriverHarness{
		client:       client,
		server:       server,
		DriverPlugin: dClient,
		logger:       logger,
		t:            t,
		impl:         d,
	}

	raw, err = client.Dispense("logmon")
	if err != nil {
		t.Fatalf("err dispensing plugin: %v", err)
	}

	h.lm = raw.(logmon.LogMon)
	return h
}

func (h *DriverHarness) Kill() {
	h.client.Close()
	h.server.Stop()
}

// MkAllocDir creates a tempory directory and allocdir structure.
// If enableLogs is set to true a logmon instance will be started to write logs
// to the LogDir of the task
// A cleanup func is returned and should be defered so as to not leak dirs
// between tests.
func (h *DriverHarness) MkAllocDir(t *TaskConfig, enableLogs bool) func() {
	dir, err := ioutil.TempDir("", "nomad_driver_harness-")
	require.NoError(h.t, err)
	t.AllocDir = dir

	allocDir := allocdir.NewAllocDir(h.logger, dir)
	require.NoError(h.t, allocDir.Build())
	taskDir := allocDir.NewTaskDir(t.Name)

	caps, err := h.Capabilities()
	require.NoError(h.t, err)

	var entries map[string]string
	fsi := caps.FSIsolation
	if fsi == cstructs.FSIsolationChroot {
		entries = config.DefaultChrootEnv
	}
	require.NoError(h.t, taskDir.Build(false, entries, fsi))

	task := &structs.Task{
		Name:      t.Name,
		Env:       t.Env,
		Resources: t.Resources.NomadResources,
	}
	taskBuilder := env.NewBuilder(mock.Node(), mock.Alloc(), task, "global")
	utils.SetEnvvars(taskBuilder, fsi, taskDir, config.DefaultConfig())

	taskEnv := taskBuilder.Build()
	if t.Env == nil {
		t.Env = taskEnv.All()
	} else {
		for k, v := range taskEnv.All() {<commit_msg>Fix panic in test setup when task does not have resources<commit_after>package drivers

import (
	"context"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"runtime"
	"time"

	hclog "github.com/hashicorp/go-hclog"
	plugin "github.com/hashicorp/go-plugin"
	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/driver/env"
	"github.com/hashicorp/nomad/client/logmon"
	cstructs "github.com/hashicorp/nomad/client/structs"
	"github.com/hashicorp/nomad/helper/testlog"
	"github.com/hashicorp/nomad/helper/uuid"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/plugins/base"
	"github.com/hashicorp/nomad/plugins/drivers/utils"
	"github.com/hashicorp/nomad/plugins/shared/hclspec"
	"github.com/mitchellh/go-testing-interface"
	"github.com/stretchr/testify/require"
)

type DriverHarness struct {
	DriverPlugin
	client *plugin.GRPCClient
	server *plugin.GRPCServer
	t      testing.T
	lm     logmon.LogMon
	logger hclog.Logger
	impl   DriverPlugin
}

func (d *DriverHarness) Impl() DriverPlugin {
	return d.impl
}
func NewDriverHarness(t testing.T, d DriverPlugin) *DriverHarness {
	logger := testlog.HCLogger(t).Named("driver_harness")
	client, server := plugin.TestPluginGRPCConn(t,
		map[string]plugin.Plugin{
			base.PluginTypeDriver: &PluginDriver{
				impl:   d,
				logger: logger.Named("driver_plugin"),
			},
			base.PluginTypeBase: &base.PluginBase{Impl: d},
			"logmon":            logmon.NewPlugin(logmon.NewLogMon(logger.Named("logmon"))),
		},
	)

	raw, err := client.Dispense(base.PluginTypeDriver)
	if err != nil {
		t.Fatalf("err dispensing plugin: %v", err)
	}

	dClient := raw.(DriverPlugin)
	h := &DriverHarness{
		client:       client,
		server:       server,
		DriverPlugin: dClient,
		logger:       logger,
		t:            t,
		impl:         d,
	}

	raw, err = client.Dispense("logmon")
	if err != nil {
		t.Fatalf("err dispensing plugin: %v", err)
	}

	h.lm = raw.(logmon.LogMon)
	return h
}

func (h *DriverHarness) Kill() {
	h.client.Close()
	h.server.Stop()
}

// MkAllocDir creates a tempory directory and allocdir structure.
// If enableLogs is set to true a logmon instance will be started to write logs
// to the LogDir of the task
// A cleanup func is returned and should be defered so as to not leak dirs
// between tests.
func (h *DriverHarness) MkAllocDir(t *TaskConfig, enableLogs bool) func() {
	dir, err := ioutil.TempDir("", "nomad_driver_harness-")
	require.NoError(h.t, err)
	t.AllocDir = dir

	allocDir := allocdir.NewAllocDir(h.logger, dir)
	require.NoError(h.t, allocDir.Build())
	taskDir := allocDir.NewTaskDir(t.Name)

	caps, err := h.Capabilities()
	require.NoError(h.t, err)

	var entries map[string]string
	fsi := caps.FSIsolation
	if fsi == cstructs.FSIsolationChroot {
		entries = config.DefaultChrootEnv
	}
	require.NoError(h.t, taskDir.Build(false, entries, fsi))

	task := &structs.Task{
		Name: t.Name,
		Env:  t.Env,
	}
	if t.Resources != nil {
		task.Resources = t.Resources.NomadResources
	}
	taskBuilder := env.NewBuilder(mock.Node(), mock.Alloc(), task, "global")
	utils.SetEnvvars(taskBuilder, fsi, taskDir, config.DefaultConfig())

	taskEnv := taskBuilder.Build()
	if t.Env == nil {
		t.Env = taskEnv.All()
	} else {
		for k, v := range taskEnv.All() {
{'fraction_non_alphanumeric': 0.13900257692890708, 'fraction_numerical': 0.0, 'mean_word_length': 3.960902255639098, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1750924', 'n_tokens_mistral': 2547, 'n_tokens_neox': 2468, 'n_words': 493}