Minor refactor of mode resolver

Signed-off-by: Evan Lezar <elezar@nvidia.com>

internal/info/auto.go

@@ -21,22 +21,40 @@ import (
 	"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvlib/info"
 )
 
+// infoInterface provides an alias for mocking.
+//
+//go:generate moq -stub -out info-interface_mock.go . infoInterface
+type infoInterface info.Interface
+
+type resolver struct {
+	logger logger.Interface
+	info   info.Interface
+}
+
 // ResolveAutoMode determines the correct mode for the platform if set to "auto"
 func ResolveAutoMode(logger logger.Interface, mode string) (rmode string) {
+	nvinfo := info.New()
+	r := resolver{
+		logger: logger,
+		info:   nvinfo,
+	}
+	return r.resolveMode(mode)
+}
+
+// resolveMode determines the correct mode for the platform if set to "auto"
+func (r resolver) resolveMode(mode string) (rmode string) {
 	if mode != "auto" {
 		return mode
 	}
 	defer func() {
-		logger.Infof("Auto-detected mode as '%v'", rmode)
+		r.logger.Infof("Auto-detected mode as '%v'", rmode)
 	}()
 
-	nvinfo := info.New()
+	isTegra, reason := r.info.IsTegraSystem()
+	r.logger.Debugf("Is Tegra-based system? %v: %v", isTegra, reason)
 
-	isTegra, reason := nvinfo.IsTegraSystem()
+	hasNVML, reason := r.info.HasNvml()
-	logger.Debugf("Is Tegra-based system? %v: %v", isTegra, reason)
+	r.logger.Debugf("Has NVML? %v: %v", hasNVML, reason)
-
-	hasNVML, reason := nvinfo.HasNvml()
-	logger.Debugf("Has NVML? %v: %v", hasNVML, reason)
 
 	if isTegra && !hasNVML {
 		return "csv"

internal/info/auto_test.go

@@ -19,8 +19,10 @@ package info
 import (
 	"testing"
 
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/config/image"
 	testlog "github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/require"
+	"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvlib/info"
 )
 
 func TestResolveAutoMode(t *testing.T) {
@@ -30,23 +32,123 @@ func TestResolveAutoMode(t *testing.T) {
 		description  string
 		mode         string
 		expectedMode string
+		info         info.Interface
+		image        image.CUDA
 	}{
 		{
 			description:  "non-auto resolves to input",
 			mode:         "not-auto",
 			expectedMode: "not-auto",
 		},
-		// TODO: The following test is brittle in that it will break on Tegra-based systems.
+		{
-		// {
+			description: "nvml non-tegra resolves to legacy",
-		// 	description:  "auto resolves to legacy",
+			mode:        "auto",
-		// 	mode:         "auto",
+			info: &infoInterfaceMock{
-		// 	expectedMode: "legacy",
+				HasNvmlFunc: func() (bool, string) {
-		// },
+					return true, "nvml"
+				},
+				IsTegraSystemFunc: func() (bool, string) {
+					return false, "tegra"
+				},
+			},
+			expectedMode: "legacy",
+		},
+		{
+			description: "non-nvml non-tegra resolves to legacy",
+			mode:        "auto",
+			info: &infoInterfaceMock{
+				HasNvmlFunc: func() (bool, string) {
+					return false, "nvml"
+				},
+				IsTegraSystemFunc: func() (bool, string) {
+					return false, "tegra"
+				},
+			},
+			expectedMode: "legacy",
+		},
+		{
+			description: "nvml tegra resolves to legacy",
+			mode:        "auto",
+			info: &infoInterfaceMock{
+				HasNvmlFunc: func() (bool, string) {
+					return true, "nvml"
+				},
+				IsTegraSystemFunc: func() (bool, string) {
+					return true, "tegra"
+				},
+			},
+			expectedMode: "legacy",
+		},
+		{
+			description: "non-nvml tegra resolves to csv",
+			mode:        "auto",
+			info: &infoInterfaceMock{
+				HasNvmlFunc: func() (bool, string) {
+					return false, "nvml"
+				},
+				IsTegraSystemFunc: func() (bool, string) {
+					return true, "tegra"
+				},
+			},
+			expectedMode: "csv",
+		},
+		{
+			description:  "cdi devices resolves to cdi",
+			mode:         "auto",
+			expectedMode: "cdi",
+			image: image.CUDA{
+				"NVIDIA_VISIBLE_DEVICES": "nvidia.com/gpu=all",
+			},
+		},
+		{
+			description:  "multiple cdi devices resolves to cdi",
+			mode:         "auto",
+			expectedMode: "cdi",
+			image: image.CUDA{
+				"NVIDIA_VISIBLE_DEVICES": "nvidia.com/gpu=0,nvidia.com/gpu=1",
+			},
+		},
+		{
+			description: "at least one non-cdi device resolves to legacy",
+			mode:        "auto",
+			image: image.CUDA{
+				"NVIDIA_VISIBLE_DEVICES": "nvidia.com/gpu=0,0",
+			},
+			info: &infoInterfaceMock{
+				HasNvmlFunc: func() (bool, string) {
+					return true, "nvml"
+				},
+				IsTegraSystemFunc: func() (bool, string) {
+					return true, "tegra"
+				},
+			},
+			expectedMode: "legacy",
+		},
+		{
+			description: "at least one non-cdi device resolves to csv",
+			mode:        "auto",
+			image: image.CUDA{
+				"NVIDIA_VISIBLE_DEVICES": "nvidia.com/gpu=0,0",
+			},
+			info: &infoInterfaceMock{
+				HasNvmlFunc: func() (bool, string) {
+					return false, "nvml"
+				},
+				IsTegraSystemFunc: func() (bool, string) {
+					return true, "tegra"
+				},
+			},
+			expectedMode: "csv",
+		},
 	}
 
 	for _, tc := range testCases {
 		t.Run(tc.description, func(t *testing.T) {
-			mode := ResolveAutoMode(logger, tc.mode)
+			r := resolver{
+				logger: logger,
+				info:   tc.info,
+			}
+			mode := r.resolveMode(tc.mode, tc.image)
 			require.EqualValues(t, tc.expectedMode, mode)
 		})
 	}

internal/info/info-interface_mock.go

@@ -0,0 +1,153 @@
+// Code generated by moq; DO NOT EDIT.
+// github.com/matryer/moq
+
+package info
+
+import (
+	"sync"
+)
+
+// Ensure, that infoInterfaceMock does implement infoInterface.
+// If this is not the case, regenerate this file with moq.
+var _ infoInterface = &infoInterfaceMock{}
+
+// infoInterfaceMock is a mock implementation of infoInterface.
+//
+//	func TestSomethingThatUsesinfoInterface(t *testing.T) {
+//
+//		// make and configure a mocked infoInterface
+//		mockedinfoInterface := &infoInterfaceMock{
+//			HasDXCoreFunc: func() (bool, string) {
+//				panic("mock out the HasDXCore method")
+//			},
+//			HasNvmlFunc: func() (bool, string) {
+//				panic("mock out the HasNvml method")
+//			},
+//			IsTegraSystemFunc: func() (bool, string) {
+//				panic("mock out the IsTegraSystem method")
+//			},
+//		}
+//
+//		// use mockedinfoInterface in code that requires infoInterface
+//		// and then make assertions.
+//
+//	}
+type infoInterfaceMock struct {
+	// HasDXCoreFunc mocks the HasDXCore method.
+	HasDXCoreFunc func() (bool, string)
+
+	// HasNvmlFunc mocks the HasNvml method.
+	HasNvmlFunc func() (bool, string)
+
+	// IsTegraSystemFunc mocks the IsTegraSystem method.
+	IsTegraSystemFunc func() (bool, string)
+
+	// calls tracks calls to the methods.
+	calls struct {
+		// HasDXCore holds details about calls to the HasDXCore method.
+		HasDXCore []struct {
+		}
+		// HasNvml holds details about calls to the HasNvml method.
+		HasNvml []struct {
+		}
+		// IsTegraSystem holds details about calls to the IsTegraSystem method.
+		IsTegraSystem []struct {
+		}
+	}
+	lockHasDXCore     sync.RWMutex
+	lockHasNvml       sync.RWMutex
+	lockIsTegraSystem sync.RWMutex
+}
+
+// HasDXCore calls HasDXCoreFunc.
+func (mock *infoInterfaceMock) HasDXCore() (bool, string) {
+	callInfo := struct {
+	}{}
+	mock.lockHasDXCore.Lock()
+	mock.calls.HasDXCore = append(mock.calls.HasDXCore, callInfo)
+	mock.lockHasDXCore.Unlock()
+	if mock.HasDXCoreFunc == nil {
+		var (
+			bOut bool
+			sOut string
+		)
+		return bOut, sOut
+	}
+	return mock.HasDXCoreFunc()
+}
+
+// HasDXCoreCalls gets all the calls that were made to HasDXCore.
+// Check the length with:
+//
+//	len(mockedinfoInterface.HasDXCoreCalls())
+func (mock *infoInterfaceMock) HasDXCoreCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockHasDXCore.RLock()
+	calls = mock.calls.HasDXCore
+	mock.lockHasDXCore.RUnlock()
+	return calls
+}
+
+// HasNvml calls HasNvmlFunc.
+func (mock *infoInterfaceMock) HasNvml() (bool, string) {
+	callInfo := struct {
+	}{}
+	mock.lockHasNvml.Lock()
+	mock.calls.HasNvml = append(mock.calls.HasNvml, callInfo)
+	mock.lockHasNvml.Unlock()
+	if mock.HasNvmlFunc == nil {
+		var (
+			bOut bool
+			sOut string
+		)
+		return bOut, sOut
+	}
+	return mock.HasNvmlFunc()
+}
+
+// HasNvmlCalls gets all the calls that were made to HasNvml.
+// Check the length with:
+//
+//	len(mockedinfoInterface.HasNvmlCalls())
+func (mock *infoInterfaceMock) HasNvmlCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockHasNvml.RLock()
+	calls = mock.calls.HasNvml
+	mock.lockHasNvml.RUnlock()
+	return calls
+}
+
+// IsTegraSystem calls IsTegraSystemFunc.
+func (mock *infoInterfaceMock) IsTegraSystem() (bool, string) {
+	callInfo := struct {
+	}{}
+	mock.lockIsTegraSystem.Lock()
+	mock.calls.IsTegraSystem = append(mock.calls.IsTegraSystem, callInfo)
+	mock.lockIsTegraSystem.Unlock()
+	if mock.IsTegraSystemFunc == nil {
+		var (
+			bOut bool
+			sOut string
+		)
+		return bOut, sOut
+	}
+	return mock.IsTegraSystemFunc()
+}
+
+// IsTegraSystemCalls gets all the calls that were made to IsTegraSystem.
+// Check the length with:
+//
+//	len(mockedinfoInterface.IsTegraSystemCalls())
+func (mock *infoInterfaceMock) IsTegraSystemCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockIsTegraSystem.RLock()
+	calls = mock.calls.IsTegraSystem
+	mock.lockIsTegraSystem.RUnlock()
+	return calls
+}