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

     1  package taskrunner
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
     8  	"github.com/hashicorp/nomad/client/consul"
     9  	"github.com/hashicorp/nomad/helper/testlog"
    10  	"github.com/hashicorp/nomad/nomad/mock"
    11  	"github.com/stretchr/testify/require"
    12  )
    13  
    14  // Statically assert the stats hook implements the expected interfaces
    15  var _ interfaces.TaskPoststartHook = (*serviceHook)(nil)
    16  var _ interfaces.TaskExitedHook = (*serviceHook)(nil)
    17  var _ interfaces.TaskPreKillHook = (*serviceHook)(nil)
    18  var _ interfaces.TaskUpdateHook = (*serviceHook)(nil)
    19  
    20  func TestUpdate_beforePoststart(t *testing.T) {
    21  	alloc := mock.Alloc()
    22  	logger := testlog.HCLogger(t)
    23  	c := consul.NewMockConsulServiceClient(t, logger)
    24  
    25  	hook := newServiceHook(serviceHookConfig{
    26  		alloc:  alloc,
    27  		task:   alloc.LookupTask("web"),
    28  		consul: c,
    29  		logger: logger,
    30  	})
    31  	require.NoError(t, hook.Update(context.Background(), &interfaces.TaskUpdateRequest{Alloc: alloc}, &interfaces.TaskUpdateResponse{}))
    32  	require.Len(t, c.GetOps(), 0)
    33  	require.NoError(t, hook.Poststart(context.Background(), &interfaces.TaskPoststartRequest{}, &interfaces.TaskPoststartResponse{}))
    34  	require.Len(t, c.GetOps(), 1)
    35  	require.NoError(t, hook.Update(context.Background(), &interfaces.TaskUpdateRequest{Alloc: alloc}, &interfaces.TaskUpdateResponse{}))
    36  	require.Len(t, c.GetOps(), 2)
    37  
    38  	// When a task exits it could be restarted with new driver info
    39  	// so Update should again wait on Poststart.
    40  
    41  	require.NoError(t, hook.Exited(context.Background(), &interfaces.TaskExitedRequest{}, &interfaces.TaskExitedResponse{}))
    42  	require.Len(t, c.GetOps(), 4)
    43  	require.NoError(t, hook.Update(context.Background(), &interfaces.TaskUpdateRequest{Alloc: alloc}, &interfaces.TaskUpdateResponse{}))
    44  	require.Len(t, c.GetOps(), 4)
    45  	require.NoError(t, hook.Poststart(context.Background(), &interfaces.TaskPoststartRequest{}, &interfaces.TaskPoststartResponse{}))
    46  	require.Len(t, c.GetOps(), 5)
    47  	require.NoError(t, hook.Update(context.Background(), &interfaces.TaskUpdateRequest{Alloc: alloc}, &interfaces.TaskUpdateResponse{}))
    48  	require.Len(t, c.GetOps(), 6)
    49  	require.NoError(t, hook.PreKilling(context.Background(), &interfaces.TaskPreKillRequest{}, &interfaces.TaskPreKillResponse{}))
    50  	require.Len(t, c.GetOps(), 8)
    51  	require.NoError(t, hook.Update(context.Background(), &interfaces.TaskUpdateRequest{Alloc: alloc}, &interfaces.TaskUpdateResponse{}))
    52  	require.Len(t, c.GetOps(), 8)
    53  
    54  }