github.com/iqoqo/nomad@v0.11.3-0.20200911112621-d7021c74d101/client/allocrunner/taskrunner/dispatch_hook_test.go (about)

     1  package taskrunner
     2  
     3  import (
     4  	"context"
     5  	"io/ioutil"
     6  	"path/filepath"
     7  	"testing"
     8  
     9  	"github.com/golang/snappy"
    10  	"github.com/hashicorp/nomad/client/allocdir"
    11  	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
    12  	"github.com/hashicorp/nomad/helper/testlog"
    13  	"github.com/hashicorp/nomad/nomad/mock"
    14  	"github.com/hashicorp/nomad/nomad/structs"
    15  	"github.com/stretchr/testify/require"
    16  )
    17  
    18  // Statically assert the stats hook implements the expected interfaces
    19  var _ interfaces.TaskPrestartHook = (*dispatchHook)(nil)
    20  
    21  // TestTaskRunner_DispatchHook_NoPayload asserts that the hook is a noop and is
    22  // marked as done if there is no dispatch payload.
    23  func TestTaskRunner_DispatchHook_NoPayload(t *testing.T) {
    24  	t.Parallel()
    25  
    26  	require := require.New(t)
    27  	ctx := context.Background()
    28  	logger := testlog.HCLogger(t)
    29  	allocDir := allocdir.NewAllocDir(logger, "nomadtest_nopayload")
    30  	defer allocDir.Destroy()
    31  
    32  	// Default mock alloc/job is not a dispatch job
    33  	alloc := mock.BatchAlloc()
    34  	task := alloc.Job.TaskGroups[0].Tasks[0]
    35  	taskDir := allocDir.NewTaskDir(task.Name)
    36  	require.NoError(taskDir.Build(false, nil))
    37  
    38  	h := newDispatchHook(alloc, logger)
    39  
    40  	req := interfaces.TaskPrestartRequest{
    41  		Task:    task,
    42  		TaskDir: taskDir,
    43  	}
    44  	resp := interfaces.TaskPrestartResponse{}
    45  
    46  	// Assert no error and Done=true as this job has no payload
    47  	require.NoError(h.Prestart(ctx, &req, &resp))
    48  	require.True(resp.Done)
    49  
    50  	// Assert payload directory is empty
    51  	files, err := ioutil.ReadDir(req.TaskDir.LocalDir)
    52  	require.NoError(err)
    53  	require.Empty(files)
    54  }
    55  
    56  // TestTaskRunner_DispatchHook_Ok asserts that dispatch payloads are written to
    57  // a file in the task dir.
    58  func TestTaskRunner_DispatchHook_Ok(t *testing.T) {
    59  	t.Parallel()
    60  
    61  	require := require.New(t)
    62  	ctx := context.Background()
    63  	logger := testlog.HCLogger(t)
    64  	allocDir := allocdir.NewAllocDir(logger, "nomadtest_dispatchok")
    65  	defer allocDir.Destroy()
    66  
    67  	// Default mock alloc/job is not a dispatch job; update it
    68  	alloc := mock.BatchAlloc()
    69  	alloc.Job.ParameterizedJob = &structs.ParameterizedJobConfig{
    70  		Payload: structs.DispatchPayloadRequired,
    71  	}
    72  	expected := []byte("hello world")
    73  	alloc.Job.Payload = snappy.Encode(nil, expected)
    74  
    75  	// Set the filename and create the task dir
    76  	task := alloc.Job.TaskGroups[0].Tasks[0]
    77  	task.DispatchPayload = &structs.DispatchPayloadConfig{
    78  		File: "out",
    79  	}
    80  	taskDir := allocDir.NewTaskDir(task.Name)
    81  	require.NoError(taskDir.Build(false, nil))
    82  
    83  	h := newDispatchHook(alloc, logger)
    84  
    85  	req := interfaces.TaskPrestartRequest{
    86  		Task:    task,
    87  		TaskDir: taskDir,
    88  	}
    89  	resp := interfaces.TaskPrestartResponse{}
    90  	require.NoError(h.Prestart(ctx, &req, &resp))
    91  	require.True(resp.Done)
    92  
    93  	filename := filepath.Join(req.TaskDir.LocalDir, task.DispatchPayload.File)
    94  	result, err := ioutil.ReadFile(filename)
    95  	require.NoError(err)
    96  	require.Equal(expected, result)
    97  }
    98  
    99  // TestTaskRunner_DispatchHook_Error asserts that on an error dispatch payloads
   100  // are not written and Done=false.
   101  func TestTaskRunner_DispatchHook_Error(t *testing.T) {
   102  	t.Parallel()
   103  
   104  	require := require.New(t)
   105  	ctx := context.Background()
   106  	logger := testlog.HCLogger(t)
   107  	allocDir := allocdir.NewAllocDir(logger, "nomadtest_dispatcherr")
   108  	defer allocDir.Destroy()
   109  
   110  	// Default mock alloc/job is not a dispatch job; update it
   111  	alloc := mock.BatchAlloc()
   112  	alloc.Job.ParameterizedJob = &structs.ParameterizedJobConfig{
   113  		Payload: structs.DispatchPayloadRequired,
   114  	}
   115  
   116  	// Cause an error by not snappy encoding the payload
   117  	alloc.Job.Payload = []byte("hello world")
   118  
   119  	// Set the filename and create the task dir
   120  	task := alloc.Job.TaskGroups[0].Tasks[0]
   121  	task.DispatchPayload = &structs.DispatchPayloadConfig{
   122  		File: "out",
   123  	}
   124  	taskDir := allocDir.NewTaskDir(task.Name)
   125  	require.NoError(taskDir.Build(false, nil))
   126  
   127  	h := newDispatchHook(alloc, logger)
   128  
   129  	req := interfaces.TaskPrestartRequest{
   130  		Task:    task,
   131  		TaskDir: taskDir,
   132  	}
   133  	resp := interfaces.TaskPrestartResponse{}
   134  
   135  	// Assert an error was returned and Done=false
   136  	require.Error(h.Prestart(ctx, &req, &resp))
   137  	require.False(resp.Done)
   138  
   139  	// Assert payload directory is empty
   140  	files, err := ioutil.ReadDir(req.TaskDir.LocalDir)
   141  	require.NoError(err)
   142  	require.Empty(files)
   143  }