src/pkg/gpubuild/gpubuild_test.go

// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gpubuild

import (
	"os"
	"path/filepath"
	"sort"
	"testing"

	"github.com/google/go-cmp/cmp"
)

func TestCopyPackageFiles(t *testing.T) {
	installerDir := t.TempDir()
	driverDir := t.TempDir()
	firmwareDir := t.TempDir()

	files := []string{
		"kernel-open/nvidia.ko",
		"kernel-open/nvidia-uvm.ko",
		"firmware/gsp_tu10x.bin",
		"firmware/gsp_ga10x.bin",
		"firmware/ucodes_tu10x.bin",
		"firmware/ucodes_ga10x.bin",
		"firmware/other_firmware.bin",
		"kernel-open/nvidia.o",
		"README.txt",
	}

	for _, relPath := range files {
		fullPath := filepath.Join(installerDir, relPath)
		if err := os.MkdirAll(filepath.Dir(fullPath), 0755); err != nil {
			t.Fatalf("os.MkdirAll(%q) failed: %v", filepath.Dir(fullPath), err)
		}
		if err := os.WriteFile(fullPath, nil, 0644); err != nil {
			t.Fatalf("os.WriteFile(%q) failed: %v", fullPath, err)
		}
	}

	if err := copyPackageFiles(installerDir, driverDir, firmwareDir); err != nil {
		t.Fatalf("copyPackageFiles(%q, %q, %q) failed: %v", installerDir, driverDir, firmwareDir, err)
	}

	driverEntries, err := os.ReadDir(driverDir)
	if err != nil {
		t.Fatalf("os.ReadDir(%q) failed: %v", driverDir, err)
	}
	var gotDrivers []string
	for _, entry := range driverEntries {
		gotDrivers = append(gotDrivers, entry.Name())
	}
	sort.Strings(gotDrivers)
	wantDrivers := []string{"nvidia-uvm.ko", "nvidia.ko"}
	if diff := cmp.Diff(wantDrivers, gotDrivers); diff != "" {
		t.Errorf("copyPackageFiles() driver files mismatch (-want +got):\n%s", diff)
	}

	firmwareEntries, err := os.ReadDir(firmwareDir)
	if err != nil {
		t.Fatalf("os.ReadDir(%q) failed: %v", firmwareDir, err)
	}
	var gotFirmware []string
	for _, entry := range firmwareEntries {
		gotFirmware = append(gotFirmware, entry.Name())
	}
	sort.Strings(gotFirmware)
	wantFirmware := []string{
		"gsp_ga10x.bin",
		"gsp_tu10x.bin",
		"ucodes_ga10x.bin",
		"ucodes_tu10x.bin",
	}
	if diff := cmp.Diff(wantFirmware, gotFirmware); diff != "" {
		t.Errorf("copyPackageFiles() firmware files mismatch (-want +got):\n%s", diff)
	}
}