github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/client/allocrunner/network_hook_test.go (about)

     1  package allocrunner
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/hashicorp/nomad/ci"
     7  	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
     8  	"github.com/hashicorp/nomad/client/taskenv"
     9  	"github.com/hashicorp/nomad/helper/testlog"
    10  	"github.com/hashicorp/nomad/nomad/mock"
    11  	"github.com/hashicorp/nomad/nomad/structs"
    12  	"github.com/hashicorp/nomad/plugins/drivers"
    13  	"github.com/hashicorp/nomad/plugins/drivers/testutils"
    14  	"github.com/stretchr/testify/require"
    15  )
    16  
    17  // statically assert network hook implements the expected interfaces
    18  var _ interfaces.RunnerPrerunHook = (*networkHook)(nil)
    19  var _ interfaces.RunnerPostrunHook = (*networkHook)(nil)
    20  
    21  type mockNetworkIsolationSetter struct {
    22  	t            *testing.T
    23  	expectedSpec *drivers.NetworkIsolationSpec
    24  	called       bool
    25  }
    26  
    27  func (m *mockNetworkIsolationSetter) SetNetworkIsolation(spec *drivers.NetworkIsolationSpec) {
    28  	m.called = true
    29  	require.Exactly(m.t, m.expectedSpec, spec)
    30  }
    31  
    32  type mockNetworkStatusSetter struct {
    33  	t              *testing.T
    34  	expectedStatus *structs.AllocNetworkStatus
    35  	called         bool
    36  }
    37  
    38  func (m *mockNetworkStatusSetter) SetNetworkStatus(status *structs.AllocNetworkStatus) {
    39  	m.called = true
    40  	require.Exactly(m.t, m.expectedStatus, status)
    41  }
    42  
    43  // Test that the prerun and postrun hooks call the setter with the expected spec when
    44  // the network mode is not host
    45  func TestNetworkHook_Prerun_Postrun(t *testing.T) {
    46  	ci.Parallel(t)
    47  
    48  	alloc := mock.Alloc()
    49  	alloc.Job.TaskGroups[0].Networks = []*structs.NetworkResource{
    50  		{
    51  			Mode: "bridge",
    52  		},
    53  	}
    54  	spec := &drivers.NetworkIsolationSpec{
    55  		Mode:   drivers.NetIsolationModeGroup,
    56  		Path:   "test",
    57  		Labels: map[string]string{"abc": "123"},
    58  	}
    59  
    60  	destroyCalled := false
    61  	nm := &testutils.MockDriver{
    62  		MockNetworkManager: testutils.MockNetworkManager{
    63  			CreateNetworkF: func(allocID string, req *drivers.NetworkCreateRequest) (*drivers.NetworkIsolationSpec, bool, error) {
    64  				require.Equal(t, alloc.ID, allocID)
    65  				return spec, false, nil
    66  			},
    67  
    68  			DestroyNetworkF: func(allocID string, netSpec *drivers.NetworkIsolationSpec) error {
    69  				destroyCalled = true
    70  				require.Equal(t, alloc.ID, allocID)
    71  				require.Exactly(t, spec, netSpec)
    72  				return nil
    73  			},
    74  		},
    75  	}
    76  	setter := &mockNetworkIsolationSetter{
    77  		t:            t,
    78  		expectedSpec: spec,
    79  	}
    80  	statusSetter := &mockNetworkStatusSetter{
    81  		t:              t,
    82  		expectedStatus: nil,
    83  	}
    84  	require := require.New(t)
    85  
    86  	envBuilder := taskenv.NewBuilder(mock.Node(), alloc, nil, alloc.Job.Region)
    87  
    88  	logger := testlog.HCLogger(t)
    89  	hook := newNetworkHook(logger, setter, alloc, nm, &hostNetworkConfigurator{}, statusSetter, envBuilder.Build())
    90  	require.NoError(hook.Prerun())
    91  	require.True(setter.called)
    92  	require.False(destroyCalled)
    93  	require.NoError(hook.Postrun())
    94  	require.True(destroyCalled)
    95  
    96  	// reset and use host network mode
    97  	setter.called = false
    98  	destroyCalled = false
    99  	alloc.Job.TaskGroups[0].Networks[0].Mode = "host"
   100  	hook = newNetworkHook(logger, setter, alloc, nm, &hostNetworkConfigurator{}, statusSetter, envBuilder.Build())
   101  	require.NoError(hook.Prerun())
   102  	require.False(setter.called)
   103  	require.False(destroyCalled)
   104  	require.NoError(hook.Postrun())
   105  	require.False(destroyCalled)
   106  }