| package installer |
| |
| import ( |
| "bytes" |
| "context" |
| "fmt" |
| "os" |
| "path" |
| "path/filepath" |
| "testing" |
| |
| "cos.googlesource.com/cos/tools.git/src/pkg/cos" |
| "cos.googlesource.com/cos/tools.git/src/pkg/fakes" |
| "cos.googlesource.com/cos/tools.git/src/pkg/gpuconfig/pb" |
| "github.com/golang/protobuf/proto" |
| ) |
| |
| func TestGetInstallerDownloadLocation(t *testing.T) { |
| for _, tc := range []struct { |
| testName string |
| metadataZone string |
| expectedLocation string |
| }{ |
| { |
| "us-west1-b", |
| "projects/123456789/zones/us-west1-b", |
| "us", |
| }, |
| { |
| "asia-east1-a", |
| "projects/123456789/zones/asia-east1-a", |
| "asia", |
| }, |
| { |
| "europe-west1-b", |
| "projects/123456789/zones/europe-west1-b", |
| "eu", |
| }, |
| { |
| "australia-southeast1-a", |
| "projects/123456789/zones/australia-southeast1-a", |
| "us", |
| }, |
| } { |
| location := getInstallerDownloadLocation(tc.metadataZone) |
| if location != tc.expectedLocation { |
| t.Errorf("%s: expect location: %s, got: %s", tc.testName, tc.expectedLocation, location) |
| } |
| } |
| } |
| |
| func TestGetPrecompiledInstallerURL(t *testing.T) { |
| ret := getPrecompiledInstallerURL("418.116.00", "73", "11647.415.0", "us") |
| expectedRet := "https://storage.googleapis.com/nvidia-drivers-us-public/nvidia-cos-project/73/tesla/418_00/418.116.00/NVIDIA-Linux-x86_64-418.116.00_73-11647-415-0.cos" |
| if ret != expectedRet { |
| t.Errorf("Unexpected return, want: %s, got: %s", expectedRet, ret) |
| } |
| } |
| |
| func TestGetGenericDriverInstallerGCSVariables(t *testing.T) { |
| |
| var testCases = []struct { |
| driverVersion string |
| bucketNvidia string |
| prefixNvidia string |
| expectedBucket string |
| expectedPath string |
| }{ |
| { |
| "535.125.09", |
| "", |
| "", |
| "nvidia-drivers-us-public", |
| "tesla/535.125.09/NVIDIA-Linux-x86_64-535.125.09.run", |
| }, |
| { |
| "535.125.09", |
| "cos-nvidia-bucket", |
| "tesla/temp", |
| "cos-nvidia-bucket", |
| "tesla/temp/NVIDIA-Linux-x86_64-535.125.09.run", |
| }, |
| { |
| "470.129.23", |
| "", |
| "tesla/temp", |
| "nvidia-drivers-us-public", |
| "tesla/470.129.23/NVIDIA-Linux-x86_64-470.129.23.run", |
| }, |
| { |
| "470.129.23", |
| "cos-nvidia-bucket", |
| "", |
| "cos-nvidia-bucket", |
| "NVIDIA-Linux-x86_64-470.129.23.run", |
| }, |
| } |
| for index, tc := range testCases { |
| t.Run(fmt.Sprintf("Test %d: TestGetGenericDriverInstallerGCSVariables", index), func(t *testing.T) { |
| actualBucket, actualPath, err := getGenericDriverInstallerGCSVariables(tc.driverVersion, tc.bucketNvidia, tc.prefixNvidia) |
| if err != nil { |
| t.Errorf("Unexpected err, want: nil, got: %v", err) |
| } |
| if actualBucket != tc.expectedBucket { |
| t.Errorf("Unexpected bucket result, want: %s, got: %s", tc.expectedBucket, actualBucket) |
| } |
| if actualPath != tc.expectedPath { |
| t.Errorf("Unexpected aritifact path result, want: %s, got: %s", tc.expectedPath, actualPath) |
| } |
| }) |
| } |
| } |
| |
| func TestDownloadGenericDriverInstaller(t *testing.T) { |
| fakeGCS := fakes.GCSForTest(t) |
| fakeBucket := "cos-tools" |
| fakePrefix := "10000.00.00/lakitu" |
| fakeGCSClient := fakeGCS.Client |
| ctx := context.Background() |
| tempDir, err := os.MkdirTemp("", "installerDir") |
| if err != nil { |
| t.Fatalf("Failed to create tempdir: %v", err) |
| } |
| defer os.RemoveAll(tempDir) |
| var mockData = []struct { |
| bucket string |
| prefix string |
| objectName string |
| content string |
| }{ |
| { |
| "nvidia-drivers-us-public", |
| "tesla/550.121.43", |
| "NVIDIA-Linux-x86_64-550.121.43.run", |
| "This is 550.121.43 run file content", |
| }, |
| { |
| "cos-tools", |
| "10000.00.00/lakitu", |
| "NVIDIA-Linux-x86_64-535.125.43.run", |
| "This is 535.125.43 run file content", |
| }, |
| { |
| "testBucket", |
| "testPrefix", |
| "NVIDIA-Linux-x86_64-470.121.43.run", |
| "This is 470.121.43 run file content", |
| }, |
| { |
| "testBucket1", |
| "", |
| "NVIDIA-Linux-x86_64-525.121.43.run", |
| "This is 525.121.43 run file content", |
| }, |
| } |
| for _, testData := range mockData { |
| fakeGCS.Objects[path.Join("/", testData.bucket, testData.prefix, testData.objectName)] = []byte(testData.content) |
| } |
| var testCases = []struct { |
| driverVersion string |
| bucketNvidia string |
| prefixNvidia string |
| expectedContent string |
| expectedPath string |
| }{ |
| { |
| "535.125.43", |
| "cos-tools", |
| "10000.00.00/lakitu", |
| "This is 535.125.43 run file content", |
| "NVIDIA-Linux-x86_64-535.125.43.run", |
| }, |
| { |
| "470.121.43", |
| "testBucket", |
| "testPrefix", |
| "This is 470.121.43 run file content", |
| "NVIDIA-Linux-x86_64-470.121.43.run", |
| }, |
| { |
| "525.121.43", |
| "testBucket1", |
| "", |
| "This is 525.121.43 run file content", |
| "NVIDIA-Linux-x86_64-525.121.43.run", |
| }, |
| { |
| "550.121.43", |
| "", |
| "", |
| "This is 550.121.43 run file content", |
| "NVIDIA-Linux-x86_64-550.121.43.run", |
| }, |
| { |
| "550.121.43", |
| "", |
| "invalidPrefix", |
| "This is 550.121.43 run file content", |
| "NVIDIA-Linux-x86_64-550.121.43.run", |
| }, |
| } |
| var FakeGCSDownloader = cos.NewGCSDownloader(fakeGCSClient, nil, fakeBucket, fakePrefix) |
| for index, tc := range testCases { |
| t.Run(fmt.Sprintf("Test %d: TestDownloadGenericDriverInstaller", index), func(t *testing.T) { |
| pathName, err := DownloadGenericDriverInstaller(ctx, FakeGCSDownloader, tc.driverVersion, tc.bucketNvidia, tc.prefixNvidia, tempDir) |
| if err != nil { |
| t.Errorf("Unexpected err, want: nil, got: %v", err) |
| } |
| if pathName != tc.expectedPath { |
| t.Errorf("Unexpected path result, want: %s, got: %s", tc.expectedPath, pathName) |
| } |
| runFileContent, err := os.ReadFile(filepath.Join(tempDir, pathName)) |
| if err != nil { |
| t.Errorf("Unexpected err, want: nil, got: %v", err) |
| } |
| if string(runFileContent) != tc.expectedContent { |
| t.Errorf("Unexpected content, want: %s, got: %s", tc.expectedContent, string(runFileContent)) |
| } |
| }) |
| } |
| |
| } |
| |
| func TestDownloadGPUDriverVersionsProto(t *testing.T) { |
| fakeGCS := fakes.GCSForTest(t) |
| fakeBucket := "cos-tools" |
| fakePrefix := "10000.00.00/lakitu" |
| fakeGCSClient := fakeGCS.Client |
| ctx := context.Background() |
| var gpuDriverProtoBin = "gpu_driver_versions.bin" |
| var mockData = &pb.GPUDriverVersionInfoList{ |
| GpuDriverVersionInfo: []*pb.GPUDriverVersionInfo{ |
| { |
| GpuType: "NVIDIA_TESLA_V100", |
| SupportedDriverVersions: []*pb.DriverVersion{ |
| { |
| Label: "DEFAULT", |
| Version: "535.154.05", |
| }, |
| { |
| Label: "LATEST", |
| Version: "535.154.05", |
| }, |
| { |
| Label: "R535", |
| Version: "535.154.05", |
| }, |
| { |
| Version: "535.129.03", |
| }, |
| { |
| Version: "535.104.12", |
| }, |
| { |
| Version: "535.104.05", |
| }, |
| { |
| Label: "R470", |
| Version: "470.223.02", |
| }, |
| { |
| Version: "470.199.02", |
| }, |
| }, |
| }, |
| }, |
| } |
| binaryMockData, err := proto.Marshal(mockData) |
| if err != nil { |
| t.Fatalf("Failed to marshal mockdata to binary array: %v", err) |
| } |
| fakeGCS.Objects[path.Join("/", fakeBucket, fakePrefix, gpuDriverProtoBin)] = binaryMockData |
| tempDir, err := os.MkdirTemp("", "mockGpuInstallDir") |
| if err != nil { |
| t.Fatalf("Failed to create tempdir: %v", err) |
| } |
| defer os.RemoveAll(tempDir) |
| expectedFilePath := filepath.Join(tempDir, gpuDriverProtoBin) |
| var FakeGCSDownloader = cos.NewGCSDownloader(fakeGCSClient, nil, fakeBucket, fakePrefix) |
| actualData, err := DownloadGPUDriverVersionsProto(ctx, FakeGCSDownloader, tempDir) |
| if err != nil { |
| t.Fatalf("DownloadGPUDriverVersionsProto returned an error: %v", err) |
| } |
| _, err = os.Stat(expectedFilePath) |
| if os.IsNotExist(err) { |
| t.Errorf("Expected file %s does not exist.", expectedFilePath) |
| } else if err != nil { |
| t.Fatalf("Error occurs when checking the existence of the file: %s", expectedFilePath) |
| } |
| |
| if !bytes.Equal(binaryMockData, actualData) { |
| t.Errorf("Data mismatch. Expected %v, got %v", binaryMockData, actualData) |
| } |
| |
| } |
