src/cmd/cos_gpu_installer/internal/installer/installer_test.go

package installer

import (
	"bytes"
	"context"
	"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 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)
	}

}