1717package modifier
1818
1919import (
20- "fmt"
2120 "testing"
2221
2322 "github.com/NVIDIA/nvidia-container-toolkit/internal/config"
24- "github.com/NVIDIA/nvidia-container-toolkit/internal/oci "
23+ "github.com/NVIDIA/nvidia-container-toolkit/internal/config/image "
2524 "github.com/opencontainers/runtime-spec/specs-go"
2625 testlog "github.com/sirupsen/logrus/hooks/test"
2726 "github.com/stretchr/testify/require"
@@ -31,54 +30,32 @@ func TestNewCSVModifier(t *testing.T) {
3130 logger , _ := testlog .NewNullLogger ()
3231
3332 testCases := []struct {
34- description string
35- cfg * config.Config
36- spec oci.Spec
37- visibleDevices string
38- expectedError error
39- expectedNil bool
33+ description string
34+ cfg * config.Config
35+ image image.CUDA
36+ expectedError error
37+ expectedNil bool
4038 }{
4139 {
42- description : "spec load error returns error" ,
43- spec : & oci.SpecMock {
44- LoadFunc : func () (* specs.Spec , error ) {
45- return nil , fmt .Errorf ("load failed" )
46- },
47- },
48- expectedError : fmt .Errorf ("load failed" ),
40+ description : "visible devices not set returns nil" ,
41+ image : image.CUDA {},
42+ expectedNil : true ,
4943 },
5044 {
51- description : "visible devices not set returns nil" ,
52- visibleDevices : "NOT_SET" ,
53- expectedNil : true ,
45+ description : "visible devices empty returns nil" ,
46+ image : image. CUDA { "NVIDIA_VISIBLE_DEVICES" : "" } ,
47+ expectedNil : true ,
5448 },
5549 {
56- description : "visible devices empty returns nil" ,
57- visibleDevices : "" ,
58- expectedNil : true ,
59- },
60- {
61- description : "visible devices 'void' returns nil" ,
62- visibleDevices : "void" ,
63- expectedNil : true ,
50+ description : "visible devices 'void' returns nil" ,
51+ image : image.CUDA {"NVIDIA_VISIBLE_DEVICES" : "void" },
52+ expectedNil : true ,
6453 },
6554 }
6655
6756 for _ , tc := range testCases {
6857 t .Run (tc .description , func (t * testing.T ) {
69- spec := tc .spec
70- if spec == nil {
71- spec = & oci.SpecMock {
72- LookupEnvFunc : func (s string ) (string , bool ) {
73- if tc .visibleDevices != "NOT_SET" && s == visibleDevicesEnvvar {
74- return tc .visibleDevices , true
75- }
76- return "" , false
77- },
78- }
79- }
80-
81- m , err := NewCSVModifier (logger , tc .cfg , spec )
58+ m , err := NewCSVModifier (logger , tc .cfg , tc .image )
8259 if tc .expectedError != nil {
8360 require .Error (t , err )
8461 } else {
0 commit comments