From f3e733f99bb273826549fd0a21f3bb1b2ee70636 Mon Sep 17 00:00:00 2001
From: Tariq Ibrahim <tibrahim@nvidia.com>
Date: Fri, 2 Feb 2024 12:33:11 -0800
Subject: [PATCH] add unit test to replicate error nvidia device rename

Signed-off-by: Tariq Ibrahim <tibrahim@nvidia.com>
---
 internal/system/nvdevices/devices_test.go | 31 +++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/internal/system/nvdevices/devices_test.go b/internal/system/nvdevices/devices_test.go
index 388f91dc..0433c389 100644
--- a/internal/system/nvdevices/devices_test.go
+++ b/internal/system/nvdevices/devices_test.go
@@ -39,6 +39,18 @@ func TestCreateControlDevices(t *testing.T) {
 		},
 	}
 
+	// nvidiaDevices550 represents the device map from the nvidia gpu drivers >= 550.40.x
+	nvidiaDevices550 := &devices.DevicesMock{
+		GetFunc: func(name devices.Name) (devices.Major, bool) {
+			devices := map[devices.Name]devices.Major{
+				"nvidia":     195,
+				"nvidia-uvm": 243,
+			}
+			d, ok := devices[name]
+			return d, ok
+		},
+	}
+
 	mknodeError := errors.New("mknode error")
 
 	testCases := []struct {
@@ -46,6 +58,7 @@ func TestCreateControlDevices(t *testing.T) {
 		root          string
 		devices       devices.Devices
 		mknodeError   error
+		hasError      bool
 		expectedError error
 		expectedCalls []struct {
 			S  string
@@ -58,6 +71,7 @@ func TestCreateControlDevices(t *testing.T) {
 			root:        "",
 			devices:     nvidiaDevices,
 			mknodeError: nil,
+			hasError:    false,
 			expectedCalls: []struct {
 				S  string
 				N1 int
@@ -73,6 +87,7 @@ func TestCreateControlDevices(t *testing.T) {
 			description: "some root specified",
 			root:        "/some/root",
 			devices:     nvidiaDevices,
+			hasError:    false,
 			mknodeError: nil,
 			expectedCalls: []struct {
 				S  string
@@ -88,6 +103,7 @@ func TestCreateControlDevices(t *testing.T) {
 		{
 			description:   "mknod error returns error",
 			devices:       nvidiaDevices,
+			hasError:      true,
 			mknodeError:   mknodeError,
 			expectedError: mknodeError,
 			// We expect the first call to this to fail, and the rest to be skipped
@@ -106,8 +122,16 @@ func TestCreateControlDevices(t *testing.T) {
 					return 0, false
 				},
 			},
+			hasError:      true,
 			expectedError: errInvalidDeviceNode,
 		},
+		{
+			description:   "nvidia device renamed from nvidia-frontend to nvidia",
+			devices:       nvidiaDevices550,
+			hasError:      true,
+			expectedError: errors.New("failed to create device node nvidiactl"),
+			expectedCalls: nil,
+		},
 	}
 
 	for _, tc := range testCases {
@@ -126,9 +150,12 @@ func TestCreateControlDevices(t *testing.T) {
 			d.mknoder = mknode
 
 			err := d.CreateNVIDIAControlDevices()
-			require.ErrorIs(t, err, tc.expectedError)
+			if tc.hasError {
+				require.ErrorContains(t, err, tc.expectedError.Error())
+			} else {
+				require.Nil(t, err)
+			}
 			require.EqualValues(t, tc.expectedCalls, mknode.MknodeCalls())
 		})
 	}
-
 }