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