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

     1  package taskrunner
     2  
     3  import (
     4  	"context"
     5  	"io/ioutil"
     6  	"net/http"
     7  	"net/http/httptest"
     8  	"os"
     9  	"path/filepath"
    10  	"sort"
    11  	"testing"
    12  
    13  	"github.com/hashicorp/nomad/client/allocdir"
    14  	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
    15  	"github.com/hashicorp/nomad/client/taskenv"
    16  	"github.com/hashicorp/nomad/helper"
    17  	"github.com/hashicorp/nomad/helper/testlog"
    18  	"github.com/hashicorp/nomad/nomad/structs"
    19  	"github.com/stretchr/testify/require"
    20  )
    21  
    22  // Statically assert the artifact hook implements the expected interface
    23  var _ interfaces.TaskPrestartHook = (*artifactHook)(nil)
    24  
    25  type mockEmitter struct {
    26  	events []*structs.TaskEvent
    27  }
    28  
    29  func (m *mockEmitter) EmitEvent(ev *structs.TaskEvent) {
    30  	m.events = append(m.events, ev)
    31  }
    32  
    33  // TestTaskRunner_ArtifactHook_Recoverable asserts that failures to download
    34  // artifacts are a recoverable error.
    35  func TestTaskRunner_ArtifactHook_Recoverable(t *testing.T) {
    36  	t.Parallel()
    37  
    38  	me := &mockEmitter{}
    39  	artifactHook := newArtifactHook(me, testlog.HCLogger(t))
    40  
    41  	req := &interfaces.TaskPrestartRequest{
    42  		TaskEnv: taskenv.NewEmptyTaskEnv(),
    43  		TaskDir: &allocdir.TaskDir{Dir: os.TempDir()},
    44  		Task: &structs.Task{
    45  			Artifacts: []*structs.TaskArtifact{
    46  				{
    47  					GetterSource: "http://127.0.0.1:0",
    48  					GetterMode:   structs.GetterModeAny,
    49  				},
    50  			},
    51  		},
    52  	}
    53  
    54  	resp := interfaces.TaskPrestartResponse{}
    55  
    56  	err := artifactHook.Prestart(context.Background(), req, &resp)
    57  
    58  	require.False(t, resp.Done)
    59  	require.NotNil(t, err)
    60  	require.True(t, structs.IsRecoverable(err))
    61  	require.Len(t, me.events, 1)
    62  	require.Equal(t, structs.TaskDownloadingArtifacts, me.events[0].Type)
    63  }
    64  
    65  // TestTaskRunnerArtifactHook_PartialDone asserts that the artifact hook skips
    66  // already downloaded artifacts when subsequent artifacts fail and cause a
    67  // restart.
    68  func TestTaskRunner_ArtifactHook_PartialDone(t *testing.T) {
    69  	t.Parallel()
    70  
    71  	me := &mockEmitter{}
    72  	artifactHook := newArtifactHook(me, testlog.HCLogger(t))
    73  
    74  	// Create a source directory with 1 of the 2 artifacts
    75  	srcdir, err := ioutil.TempDir("", "nomadtest-src")
    76  	require.NoError(t, err)
    77  	defer func() {
    78  		require.NoError(t, os.RemoveAll(srcdir))
    79  	}()
    80  
    81  	// Only create one of the 2 artifacts to cause an error on first run.
    82  	file1 := filepath.Join(srcdir, "foo.txt")
    83  	require.NoError(t, ioutil.WriteFile(file1, []byte{'1'}, 0644))
    84  
    85  	// Test server to serve the artifacts
    86  	ts := httptest.NewServer(http.FileServer(http.Dir(srcdir)))
    87  	defer ts.Close()
    88  
    89  	// Create the target directory.
    90  	destdir, err := ioutil.TempDir("", "nomadtest-dest")
    91  	require.NoError(t, err)
    92  	defer func() {
    93  		require.NoError(t, os.RemoveAll(destdir))
    94  	}()
    95  
    96  	req := &interfaces.TaskPrestartRequest{
    97  		TaskEnv: taskenv.NewEmptyTaskEnv(),
    98  		TaskDir: &allocdir.TaskDir{Dir: destdir},
    99  		Task: &structs.Task{
   100  			Artifacts: []*structs.TaskArtifact{
   101  				{
   102  					GetterSource: ts.URL + "/foo.txt",
   103  					GetterMode:   structs.GetterModeAny,
   104  				},
   105  				{
   106  					GetterSource: ts.URL + "/bar.txt",
   107  					GetterMode:   structs.GetterModeAny,
   108  				},
   109  			},
   110  		},
   111  	}
   112  
   113  	resp := interfaces.TaskPrestartResponse{}
   114  
   115  	// On first run file1 (foo) should download but file2 (bar) should
   116  	// fail.
   117  	err = artifactHook.Prestart(context.Background(), req, &resp)
   118  
   119  	require.NotNil(t, err)
   120  	require.True(t, structs.IsRecoverable(err))
   121  	require.Len(t, resp.State, 1)
   122  	require.False(t, resp.Done)
   123  	require.Len(t, me.events, 1)
   124  	require.Equal(t, structs.TaskDownloadingArtifacts, me.events[0].Type)
   125  
   126  	// Remove file1 from the server so it errors if its downloaded again.
   127  	require.NoError(t, os.Remove(file1))
   128  
   129  	// Write file2 so artifacts can download successfully
   130  	file2 := filepath.Join(srcdir, "bar.txt")
   131  	require.NoError(t, ioutil.WriteFile(file2, []byte{'1'}, 0644))
   132  
   133  	// Mock TaskRunner by copying state from resp to req and reset resp.
   134  	req.PreviousState = helper.CopyMapStringString(resp.State)
   135  
   136  	resp = interfaces.TaskPrestartResponse{}
   137  
   138  	// Retry the download and assert it succeeds
   139  	err = artifactHook.Prestart(context.Background(), req, &resp)
   140  
   141  	require.NoError(t, err)
   142  	require.True(t, resp.Done)
   143  	require.Len(t, resp.State, 2)
   144  
   145  	// Assert both files downloaded properly
   146  	files, err := filepath.Glob(filepath.Join(destdir, "*.txt"))
   147  	require.NoError(t, err)
   148  	sort.Strings(files)
   149  	require.Contains(t, files[0], "bar.txt")
   150  	require.Contains(t, files[1], "foo.txt")
   151  
   152  	// Stop the test server entirely and assert that re-running works
   153  	ts.Close()
   154  	req.PreviousState = helper.CopyMapStringString(resp.State)
   155  	resp = interfaces.TaskPrestartResponse{}
   156  	err = artifactHook.Prestart(context.Background(), req, &resp)
   157  	require.NoError(t, err)
   158  	require.True(t, resp.Done)
   159  	require.Len(t, resp.State, 2)
   160  }