src/pkg/gpuconfig/gpu_driver_versions_test.go

package gpuconfig

import (
	"errors"
	"fmt"
	"os"
	"strings"
	"testing"

	"cos.googlesource.com/cos/tools.git/src/pkg/gpuconfig/pb"
	"github.com/golang/protobuf/proto"
	"github.com/google/go-cmp/cmp"
	"google.golang.org/protobuf/encoding/prototext"
)

const testProtoDataPath = "./testdata/gpu_driver_versions_test.textproto"

func TestGetGPUDriverVersion(t *testing.T) {
	textProtoData, err := os.ReadFile(testProtoDataPath)
	if err != nil {
		t.Errorf("Cannot read the test textproto data: %v", err)
	}
	var gpuDriverVersionInfoList = &pb.GPUDriverVersionInfoList{}
	err = (prototext.UnmarshalOptions{DiscardUnknown: true, AllowPartial: true}).Unmarshal(textProtoData, gpuDriverVersionInfoList)
	if err != nil {
		t.Errorf("failed when parsing the GPU driver version: %v", err)
	}
	testData, err := proto.Marshal(gpuDriverVersionInfoList)
	if err != nil {
		t.Errorf("fail to encode to the binary proto data: %v", err)
	}
	var testCases = []struct {
		gpuProtoContent       []byte
		gpuType               string
		input                 string
		fallback              bool
		expectedDriverVersion string
		err                   error
	}{
		{
			nil,
			"NVIDIA_L4",
			"",
			false,
			"",
			errors.New("the gpu proto content must not be empty"),
		},
		{
			testData,
			"",
			"",
			false,
			"",
			errors.New("the GPU type must not be empty"),
		},
		{
			testData,
			"   nvidia_l4 ",
			"",
			false,
			"535.154.05",
			nil,
		},
		{
			testData,
			"NVIDIA_L4",
			"560.38",
			false,
			"560.38",
			nil,
		},
		{
			testData,
			"NVIDIA_L4",
			"DEFAULT",
			false,
			"535.154.05",
			nil,
		},
		{
			testData,
			"NVIDIA_L4",
			"default",
			false,
			"535.154.05",
			nil,
		},
		{
			testData,
			"NVIDIA_L4",
			"latest",
			false,
			"535.154.05",
			nil,
		},
		{
			testData,
			"NVIDIA_L4",
			"R535",
			false,
			"535.154.05",
			nil,
		},
		{
			testData,
			"NVIDIA_L4",
			"r535",
			false,
			"535.154.05",
			nil,
		},
		{
			testData,
			"NVIDIA_L4",
			"535.129.03",
			false,
			"535.129.03",
			nil,
		},
		{
			testData,
			"NVIDIA_L4",
			"R470",
			false,
			"",
			errors.New("not supported for GPU type"),
		},
		{
			testData,
			"NVIDIA_L4",
			"R470",
			true,
			"535.154.05",
			nil,
		},
		{
			testData,
			"NVIDIA_L4",
			"R525",
			true,
			"535.154.05",
			nil,
		},
		{
			testData,
			"invalidGPU",
			"R525",
			true,
			"",
			errors.New("no supported driver versions found for gpu"),
		},
		{
			testData,
			"NVIDIA_L4",
			"invalidInput",
			false,
			"",
			errors.New("the input is invalid"),
		},
		{
			testData,
			"NVIDIA_L4",
			"invalidInput",
			true,
			"",
			errors.New("the input is invalid"),
		},
		{
			testData,
			"no_gpu",
			"r525",
			true,
			"535.161.08",
			nil,
		},
		{
			testData,
			"no_gpu",
			"r525",
			true,
			"535.161.08",
			nil,
		},
		{
			testData,
			"Sample_GPU_without_default_label",
			"r525",
			true,
			"",
			errors.New("unable to parse gpu_type"),
		},
	}
	for index, tc := range testCases {
		t.Run(fmt.Sprintf("Test %d: GetGPUDriverVersion: with gpuType: %s, input: %s, fallback flag: %v", index, tc.gpuType, tc.input, tc.fallback), func(t *testing.T) {
			driverVersion, err := GetGPUDriverVersion(tc.gpuProtoContent, tc.gpuType, tc.input, tc.fallback)
			if tc.err == nil {
				if err != nil {
					t.Errorf("Failed to GetGPUDriverVersion with error: %v", err)
				}
				if !cmp.Equal(tc.expectedDriverVersion, driverVersion) {
					t.Errorf("Test GetGPUDriverVersion failed: the expected driver version is: %s, but the actual driver version is: %s.", tc.expectedDriverVersion, driverVersion)
				}
			} else {
				if err == nil || !strings.Contains(err.Error(), tc.err.Error()) {
					t.Errorf("Test GetGPUDriverVersion failed: the error been sent out is not correct, the exected: %s, the actual: %s", tc.err.Error(), err.Error())
				}
			}
		})
	}
}
func TestGetGPUDriverVersions(t *testing.T) {
	textProtoData, err := os.ReadFile(testProtoDataPath)
	if err != nil {
		t.Errorf("Cannot read the test textproto data: %v", err)
	}
	var gpuDriverVersionInfoList = &pb.GPUDriverVersionInfoList{}
	err = (prototext.UnmarshalOptions{DiscardUnknown: true, AllowPartial: true}).Unmarshal(textProtoData, gpuDriverVersionInfoList)
	if err != nil {
		t.Errorf("failed when parsing the GPU driver version: %v", err)
	}
	testData, err := proto.Marshal(gpuDriverVersionInfoList)
	if err != nil {
		t.Errorf("fail to encode to the binary proto data: %v", err)
	}
	var testCases = []struct {
		gpuProtoContent        []byte
		gpuType                string
		expectedDriverVersions []*pb.DriverVersion
		err                    error
	}{
		{
			nil,
			"NVIDIA_L4",
			nil,
			errors.New("the gpu proto content must not be empty"),
		},
		{
			testData,
			"",
			nil,
			errors.New("the GPU type must not be empty"),
		},
		{
			testData,
			"NVIDIA_TESLA_V100",
			[]*pb.DriverVersion{
				{Version: "535.154.05", Label: "DEFAULT"},
				{Version: "535.154.05", Label: "LATEST"},
				{Version: "535.154.05", Label: "R535"},
				{Version: "535.129.03"},
				{Version: "535.104.12"},
				{Version: "535.104.05"},
				{Version: "470.223.02", Label: "R470"},
				{Version: "470.199.02"},
			},
			nil,
		},
		{
			testData,
			"NVIDIA_TESLA_P100",
			[]*pb.DriverVersion{
				{Version: "535.154.05", Label: "DEFAULT"},
				{Version: "535.154.05", Label: "LATEST"},
				{Version: "535.154.05", Label: "R535"},
				{Version: "535.129.03"},
				{Version: "535.104.12"},
				{Version: "535.104.05"},
				{Version: "470.223.02", Label: "R470"},
				{Version: "470.199.02"},
			},
			nil,
		},
		{
			testData,
			"InvalidGPU",
			nil,
			errors.New("no supported driver versions found for gpu"),
		},
	}
	for index, tc := range testCases {
		t.Run(fmt.Sprintf("Test %d: GetGPUDriverVersions: with gpuType: %s", index, tc.gpuType), func(t *testing.T) {
			driverVersions, err := GetGPUDriverVersions(tc.gpuProtoContent, tc.gpuType)
			if tc.err == nil {
				if err != nil {
					t.Errorf("Failed to GetGPUDriverVersions with error: %v", err)
				}
				if len(driverVersions) != len(tc.expectedDriverVersions) {
					t.Errorf("Test GetGPUDriverVersions failed: the length of expected gpu drivers is %v, but the actual lenght is %v", len(tc.expectedDriverVersions), len(driverVersions))
				}
				for index, driverVersion := range driverVersions {
					expectedDriverVersion := tc.expectedDriverVersions[index]
					if expectedDriverVersion.Label != driverVersion.Label {
						t.Errorf("Test GetGPUDriverVersions failed: the expected driver version list contains DriverVersion{Label: %s, Version: %s},"+
							"but the actual driver version list contains {DriverVersion{Label: %s, Version: %s}}.", expectedDriverVersion.Label, expectedDriverVersion.Version,
							driverVersion.Label, driverVersion.Version)
					}
					if expectedDriverVersion.Version != driverVersion.Version {
						t.Errorf("Test GetGPUDriverVersions failed: the expected driver version list contains DriverVersion{Label: %s, Version: %s},"+
							"but the actual driver version list contains {DriverVersion{Label: %s, Version: %s}}.", expectedDriverVersion.Label, expectedDriverVersion.Version,
							driverVersion.Label, driverVersion.Version)
					}
				}
			} else {
				if err == nil || !strings.Contains(err.Error(), tc.err.Error()) {
					t.Errorf("Test GetGPUDriverVersions failed: the error been sent out is not correct, the exected: %s, the actual: %s", tc.err.Error(), err.Error())
				}
			}
		})
	}
}

func TestGetVGPUDriverVersion(t *testing.T) {
	textProtoData, err := os.ReadFile(testProtoDataPath)
	if err != nil {
		t.Errorf("Cannot read the test textproto data: %v", err)
	}

	var testCases = []struct {
		name                  string
		gpuType               string
		input                 string
		fallback              bool
		hostDriverVersion     string
		isVGPU                bool
		expectedDriverVersion string
		err                   error
	}{
		{
			name:                  "missing host driver version",
			gpuType:               "NVIDIA_RTX_PRO_6000",
			input:                 "default",
			fallback:              false,
			hostDriverVersion:     "",
			isVGPU:                true,
			expectedDriverVersion: "",
			err:                   errors.New("host driver version missing"),
		},
		{
			name:                  "precise version matching host branch",
			gpuType:               "NVIDIA_RTX_PRO_6000",
			input:                 "575.57.08",
			fallback:              false,
			hostDriverVersion:     "575.10.10",
			isVGPU:                true,
			expectedDriverVersion: "575.57.08-grid-gcp",
			err:                   nil,
		},
		{
			name:                  "precise grid version matching host branch",
			gpuType:               "NVIDIA_RTX_PRO_6000",
			input:                 "575.57.08-grid",
			fallback:              false,
			hostDriverVersion:     "575.10.10",
			isVGPU:                true,
			expectedDriverVersion: "575.57.08-grid",
			err:                   nil,
		},
		{
			name:                  "precise grid-gcp version matching host branch",
			gpuType:               "NVIDIA_RTX_PRO_6000",
			input:                 "575.57.08-grid-gcp",
			fallback:              false,
			hostDriverVersion:     "575.10.10",
			isVGPU:                true,
			expectedDriverVersion: "575.57.08-grid-gcp",
			err:                   nil,
		},
		{
			name:                  "precise version incompatible with host, fallback disabled",
			gpuType:               "NVIDIA_RTX_PRO_6000",
			input:                 "575.00.00",
			fallback:              false,
			hostDriverVersion:     "550.10.10",
			isVGPU:                true,
			expectedDriverVersion: "",
			err:                   errors.New("is either not supported for vGPU type"),
		},
		{
			name:                  "precise version only has grid-gcp",
			gpuType:               "NVIDIA_RTX_PRO_6000",
			input:                 "550.00.00",
			fallback:              true,
			hostDriverVersion:     "550.10.10",
			isVGPU:                true,
			expectedDriverVersion: "550.00.00-grid-gcp",
			err:                   nil,
		},
		{
			name:                  "precise version only has grid",
			gpuType:               "NVIDIA_RTX_PRO_6000",
			input:                 "540.00.00",
			fallback:              false,
			hostDriverVersion:     "540.10.10",
			isVGPU:                true,
			expectedDriverVersion: "540.00.00-grid",
			err:                   nil,
		},
		{
			name:                  "label matched, incompatible with host, fallback disabled",
			gpuType:               "NVIDIA_RTX_PRO_6000",
			input:                 "R550",
			fallback:              false,
			hostDriverVersion:     "575.10.10",
			isVGPU:                true,
			expectedDriverVersion: "",
			err:                   errors.New("and fallback is disabled"),
		},
		{
			name:                  "label matched, incompatible, fallback enabled",
			gpuType:               "NVIDIA_RTX_PRO_6000",
			input:                 "R550",
			fallback:              true,
			hostDriverVersion:     "575.10.10",
			isVGPU:                true,
			expectedDriverVersion: "575.57.08-grid-gcp",
			err:                   nil,
		},
		{
			name:                  "unknown label, fallback disabled",
			gpuType:               "NVIDIA_RTX_PRO_6000",
			input:                 "UNKNOWN",
			fallback:              false,
			hostDriverVersion:     "575.10.10",
			isVGPU:                true,
			expectedDriverVersion: "",
			err:                   errors.New("is not supported for vGPU type"),
		},
		{
			name:                  "unknown label, fallback enabled",
			gpuType:               "NVIDIA_RTX_PRO_6000",
			input:                 "UNKNOWN",
			fallback:              true,
			hostDriverVersion:     "575.10.10",
			isVGPU:                true,
			expectedDriverVersion: "575.57.08-grid-gcp",
			err:                   nil,
		},
		{
			name:                  "fallback finds exact host version as grid when host driver incompatible, only grid available",
			gpuType:               "NVIDIA_RTX_PRO_6000",
			input:                 "R550", // R550 is incompatible with 540 host branch
			fallback:              true,
			hostDriverVersion:     "540.00.00",
			isVGPU:                true,
			expectedDriverVersion: "540.00.00-grid",
			err:                   nil,
		},
		{
			name:                  "fallback when no matching host branch, returns latest grid",
			gpuType:               "NVIDIA_RTX_PRO_6000",
			input:                 "default",
			fallback:              true,
			hostDriverVersion:     "999.99.99",
			isVGPU:                true,
			expectedDriverVersion: "575.57.08-grid-gcp",
			err:                   nil,
		},
		{
			name:                  "precise grid uppercase version matching host branch",
			gpuType:               "NVIDIA_RTX_PRO_6000",
			input:                 "575.57.08-GRID",
			fallback:              false,
			hostDriverVersion:     "575.10.10",
			isVGPU:                true,
			expectedDriverVersion: "575.57.08-grid",
			err:                   nil,
		},
		{
			name:                  "precise grid-gcp uppercase version matching host branch",
			gpuType:               "NVIDIA_RTX_PRO_6000",
			input:                 "575.57.08-GRID-GCP",
			fallback:              false,
			hostDriverVersion:     "575.10.10",
			isVGPU:                true,
			expectedDriverVersion: "575.57.08-grid-gcp",
			err:                   nil,
		},
	}

	var gpuDriverVersionInfoList = &pb.GPUDriverVersionInfoList{}
	err = (prototext.UnmarshalOptions{DiscardUnknown: true, AllowPartial: true}).Unmarshal(textProtoData, gpuDriverVersionInfoList)
	if err != nil {
		t.Fatalf("failed when parsing the test textproto data: %v", err)
	}
	testData, err := proto.Marshal(gpuDriverVersionInfoList)
	if err != nil {
		t.Fatalf("fail to encode to the binary proto data: %v", err)
	}

	for index, tc := range testCases {
		t.Run(fmt.Sprintf("Test %d: %s", index, tc.name), func(t *testing.T) {
			gpuDriverVersions, err := parseGPUDriverVersion(testData, tc.gpuType)
			if err != nil {
				t.Fatalf("Failed to parse GPU driver versions: %v", err)
			}
			gpuDriverVersions.hostDriverVersion = tc.hostDriverVersion
			gpuDriverVersions.gpu.IsVGPU = tc.isVGPU

			driverVersion, err := getVGPUVersion(gpuDriverVersions, tc.input, tc.fallback, gpuDriverVersions.gpu)

			if tc.err == nil {
				if err != nil {
					t.Errorf("Failed to getVGPUVersion with error: %v", err)
				}
				if driverVersion != tc.expectedDriverVersion {
					t.Errorf("Expected driver version %s, but got %s", tc.expectedDriverVersion, driverVersion)
				}
			} else {
				if err == nil || !strings.Contains(err.Error(), tc.err.Error()) {
					t.Errorf("Expected error containing '%s', but got: %v", tc.err.Error(), err)
				}
			}
		})
	}
}

func TestSortGridVersions(t *testing.T) {
	testCases := []struct {
		name     string
		input    []string
		expected []string
	}{
		{
			name:     "versions with leading zeroes sorted numerically",
			input:    []string{"580.126.09-grid", "580.95.05-grid", "580.159.03-grid-gcp", "580.82.07-grid"},
			expected: []string{"580.82.07-grid", "580.95.05-grid", "580.126.09-grid", "580.159.03-grid-gcp"},
		},
		{
			name:     "same version prioritizes grid-gcp over grid",
			input:    []string{"580.126.09-grid-gcp", "580.126.09-grid"},
			expected: []string{"580.126.09-grid", "580.126.09-grid-gcp"},
		},
		{
			name:     "case insensitivity with uppercase grid suffixes",
			input:    []string{"580.126.09-GRID", "580.95.05-GRID-GCP"},
			expected: []string{"580.95.05-grid-gcp", "580.126.09-grid"},
		},
		{
			name:     "two-part and three-part version numbers",
			input:    []string{"560.38.01-grid", "560.38-grid"},
			expected: []string{"560.38-grid", "560.38.01-grid"},
		},
		{
			name:     "filters out non-grid versions",
			input:    []string{"580.173.02", "580.126.09-grid-gcp", "580.159.04", "580.95.05-grid"},
			expected: []string{"580.95.05-grid", "580.126.09-grid-gcp"},
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			actual := sortGridVersions(tc.input)
			if diff := cmp.Diff(tc.expected, actual); diff != "" {
				t.Errorf("sortGridVersions mismatch (-expected +actual):\n%s", diff)
			}
		})
	}
}

func TestCompareNvidiaVersions(t *testing.T) {
	testCases := []struct {
		v1       string
		v2       string
		expected int
	}{
		{"580.126.09-grid", "580.95.05-grid", 1},
		{"580.95.05-grid", "580.126.09-grid", -1},
		{"580.126.09-grid", "580.126.09-grid", 0},
		{"580.126.09-grid", "580.126.09-grid-gcp", -1},
		{"580.126.09-grid-gcp", "580.126.09-grid", 1},
		{"580.159.03-grid-gcp", "580.126.09-grid-gcp", 1},
		{"580.82.07-grid", "580.95.05-grid", -1},
		{"575.57.08-GRID", "575.57.08-grid", 0},
	}

	for _, tc := range testCases {
		t.Run(fmt.Sprintf("%s_vs_%s", tc.v1, tc.v2), func(t *testing.T) {
			actual := compareNvidiaVersions(tc.v1, tc.v2)
			if actual != tc.expected {
				t.Errorf("compareNvidiaVersions(%q, %q) = %d, expected %d", tc.v1, tc.v2, actual, tc.expected)
			}
		})
	}
}