github.com/iqoqo/nomad@v0.11.3-0.20200911112621-d7021c74d101/devices/gpu/nvidia/device_test.go (about)

     1  package nvidia
     2  
     3  import (
     4  	"testing"
     5  
     6  	hclog "github.com/hashicorp/go-hclog"
     7  	"github.com/hashicorp/nomad/devices/gpu/nvidia/nvml"
     8  	"github.com/hashicorp/nomad/plugins/device"
     9  	"github.com/stretchr/testify/require"
    10  )
    11  
    12  type MockNvmlClient struct {
    13  	FingerprintError            error
    14  	FingerprintResponseReturned *nvml.FingerprintData
    15  
    16  	StatsError            error
    17  	StatsResponseReturned []*nvml.StatsData
    18  }
    19  
    20  func (c *MockNvmlClient) GetFingerprintData() (*nvml.FingerprintData, error) {
    21  	return c.FingerprintResponseReturned, c.FingerprintError
    22  }
    23  
    24  func (c *MockNvmlClient) GetStatsData() ([]*nvml.StatsData, error) {
    25  	return c.StatsResponseReturned, c.StatsError
    26  }
    27  
    28  func TestReserve(t *testing.T) {
    29  	for _, testCase := range []struct {
    30  		Name                string
    31  		ExpectedReservation *device.ContainerReservation
    32  		ExpectedError       error
    33  		Device              *NvidiaDevice
    34  		RequestedIDs        []string
    35  	}{
    36  		{
    37  			Name:                "All RequestedIDs are not managed by Device",
    38  			ExpectedReservation: nil,
    39  			ExpectedError: &reservationError{[]string{
    40  				"UUID1",
    41  				"UUID2",
    42  				"UUID3",
    43  			}},
    44  			RequestedIDs: []string{
    45  				"UUID1",
    46  				"UUID2",
    47  				"UUID3",
    48  			},
    49  			Device: &NvidiaDevice{
    50  				logger: hclog.NewNullLogger(),
    51  			},
    52  		},
    53  		{
    54  			Name:                "Some RequestedIDs are not managed by Device",
    55  			ExpectedReservation: nil,
    56  			ExpectedError: &reservationError{[]string{
    57  				"UUID1",
    58  				"UUID2",
    59  			}},
    60  			RequestedIDs: []string{
    61  				"UUID1",
    62  				"UUID2",
    63  				"UUID3",
    64  			},
    65  			Device: &NvidiaDevice{
    66  				devices: map[string]struct{}{
    67  					"UUID3": {},
    68  				},
    69  				logger: hclog.NewNullLogger(),
    70  			},
    71  		},
    72  		{
    73  			Name: "All RequestedIDs are managed by Device",
    74  			ExpectedReservation: &device.ContainerReservation{
    75  				Envs: map[string]string{
    76  					NvidiaVisibleDevices: "UUID1,UUID2,UUID3",
    77  				},
    78  			},
    79  			ExpectedError: nil,
    80  			RequestedIDs: []string{
    81  				"UUID1",
    82  				"UUID2",
    83  				"UUID3",
    84  			},
    85  			Device: &NvidiaDevice{
    86  				devices: map[string]struct{}{
    87  					"UUID1": {},
    88  					"UUID2": {},
    89  					"UUID3": {},
    90  				},
    91  				logger: hclog.NewNullLogger(),
    92  			},
    93  		},
    94  		{
    95  			Name:                "No IDs requested",
    96  			ExpectedReservation: &device.ContainerReservation{},
    97  			ExpectedError:       nil,
    98  			RequestedIDs:        nil,
    99  			Device: &NvidiaDevice{
   100  				devices: map[string]struct{}{
   101  					"UUID1": {},
   102  					"UUID2": {},
   103  					"UUID3": {},
   104  				},
   105  				logger: hclog.NewNullLogger(),
   106  			},
   107  		},
   108  	} {
   109  		actualReservation, actualError := testCase.Device.Reserve(testCase.RequestedIDs)
   110  		req := require.New(t)
   111  		req.Equal(testCase.ExpectedReservation, actualReservation)
   112  		req.Equal(testCase.ExpectedError, actualError)
   113  	}
   114  }