bee/audit/internal/platform/nvidia_settings_test.go

package platform

import (
	"os/exec"
	"strings"
	"testing"
)

// fakeNvidiaSMI stubs satExecCommand for nvidia-smi calls made by the
// GPU-settings code: the CSV query (index,...,power.default_limit), the
// per-GPU "conf-compute -q" read, and write commands (-e/-mig/-pl/-scc).
// ccStateByIndex maps GPU index -> CC State text ("ON"/"OFF"); a missing
// entry simulates conf-compute being unsupported (query fails).
func fakeNvidiaSMI(t *testing.T, queryOut string, ccStateByIndex map[int]string, writeCalls *[]string) {
	t.Helper()
	old := satExecCommand
	satExecCommand = func(name string, args ...string) *exec.Cmd {
		if name != "nvidia-smi" {
			return exec.Command(name, args...)
		}
		if len(args) > 0 && strings.HasPrefix(args[0], "--query-gpu=") {
			return exec.Command("printf", "%s", queryOut)
		}
		if len(args) > 0 && args[0] == "conf-compute" {
			// args: conf-compute -i <index> -q
			idx := -1
			for i, a := range args {
				if a == "-i" && i+1 < len(args) {
					var n int
					fscanInt(args[i+1], &n)
					idx = n
				}
			}
			if len(args) > 0 && args[len(args)-1] == "-q" {
				state, ok := ccStateByIndex[idx]
				if !ok {
					return exec.Command("false")
				}
				return exec.Command("printf", "CC State                : %s\\n", state)
			}
			if writeCalls != nil {
				*writeCalls = append(*writeCalls, strings.Join(args, " "))
			}
			return exec.Command("true")
		}
		if writeCalls != nil {
			*writeCalls = append(*writeCalls, strings.Join(args, " "))
		}
		return exec.Command("true")
	}
	t.Cleanup(func() { satExecCommand = old })
}

func fscanInt(s string, out *int) {
	n := 0
	neg := false
	for i, c := range s {
		if i == 0 && c == '-' {
			neg = true
			continue
		}
		if c < '0' || c > '9' {
			return
		}
		n = n*10 + int(c-'0')
	}
	if neg {
		n = -n
	}
	*out = n
}

func TestListNvidiaGPUSettingsParsesCSV(t *testing.T) {
	fakeNvidiaSMI(t,
		"0, H100 80GB HBM3, Enabled, Enabled, Disabled, Disabled, 350.00, 100.00, 700.00, 700.00\n"+
			"1, H100 80GB HBM3, Disabled, Enabled, N/A, N/A, 300.00, 100.00, 700.00, 700.00\n",
		map[int]string{0: "OFF", 1: "OFF"}, nil)

	s := &System{}
	got, err := s.ListNvidiaGPUSettings()
	if err != nil {
		t.Fatalf("err=%v", err)
	}
	if len(got) != 2 {
		t.Fatalf("len=%d want 2 (%#v)", len(got), got)
	}
	if got[0].ECCCurrent != "Enabled" || got[0].PowerLimitW != 350 || got[0].PowerDefaultLimitW != 700 {
		t.Fatalf("gpu0=%#v", got[0])
	}
	if got[0].MIGCurrent != "Disabled" || got[0].CCState != "OFF" {
		t.Fatalf("gpu0 mig/cc=%#v", got[0])
	}
	if got[1].ECCCurrent != "Disabled" || got[1].ECCPending != "Enabled" || got[1].MIGCurrent != "N/A" {
		t.Fatalf("gpu1=%#v", got[1])
	}
}

func TestListNvidiaGPUSettingsCCUnsupportedLeavesEmptyState(t *testing.T) {
	fakeNvidiaSMI(t,
		"0, H100, Enabled, Enabled, Disabled, Disabled, 350.00, 100.00, 700.00, 700.00\n",
		map[int]string{}, nil)

	s := &System{}
	got, err := s.ListNvidiaGPUSettings()
	if err != nil {
		t.Fatalf("err=%v", err)
	}
	if len(got) != 1 || got[0].CCState != "" {
		t.Fatalf("got=%#v want empty CCState when conf-compute query fails", got)
	}
}

func TestSetNvidiaGPUECCPassesCorrectFlag(t *testing.T) {
	var calls []string
	fakeNvidiaSMI(t, "", nil, &calls)

	s := &System{}
	if _, err := s.SetNvidiaGPUECC(2, true); err != nil {
		t.Fatalf("err=%v", err)
	}
	if len(calls) != 1 || calls[0] != "-i 2 -e 1" {
		t.Fatalf("calls=%v want [-i 2 -e 1]", calls)
	}

	if _, err := s.SetNvidiaGPUECC(0, false); err != nil {
		t.Fatalf("err=%v", err)
	}
	if calls[1] != "-i 0 -e 0" {
		t.Fatalf("calls[1]=%q want -i 0 -e 0", calls[1])
	}
}

func TestSetNvidiaGPUMIGPassesCorrectFlag(t *testing.T) {
	var calls []string
	fakeNvidiaSMI(t, "", nil, &calls)

	s := &System{}
	if _, err := s.SetNvidiaGPUMIG(3, true); err != nil {
		t.Fatalf("err=%v", err)
	}
	if len(calls) != 1 || calls[0] != "-i 3 -mig 1" {
		t.Fatalf("calls=%v want [-i 3 -mig 1]", calls)
	}
}

func TestSetNvidiaGPUCCModePassesCorrectFlag(t *testing.T) {
	var calls []string
	fakeNvidiaSMI(t, "", nil, &calls)

	s := &System{}
	if _, err := s.SetNvidiaGPUCCMode(1, false); err != nil {
		t.Fatalf("err=%v", err)
	}
	if len(calls) != 1 || calls[0] != "conf-compute -i 1 -scc 0" {
		t.Fatalf("calls=%v want [conf-compute -i 1 -scc 0]", calls)
	}
}

func TestSetNvidiaGPUPowerLimitFormatsWatts(t *testing.T) {
	var calls []string
	fakeNvidiaSMI(t, "", nil, &calls)

	s := &System{}
	if _, err := s.SetNvidiaGPUPowerLimit(1, 450); err != nil {
		t.Fatalf("err=%v", err)
	}
	if len(calls) != 1 || calls[0] != "-i 1 -pl 450" {
		t.Fatalf("calls=%v want [-i 1 -pl 450]", calls)
	}
}

func TestResetNvidiaGPUDefaultsOnlyTouchesDriftedGPUs(t *testing.T) {
	var calls []string
	fakeNvidiaSMI(t,
		"0, H100, Enabled, Enabled, Enabled, Enabled, 350.00, 100.00, 700.00, 700.00\n"+
			"1, H100, Disabled, Disabled, Disabled, Disabled, 700.00, 100.00, 700.00, 700.00\n",
		map[int]string{0: "ON", 1: "OFF"}, &calls)

	s := &System{}
	out, err := s.ResetNvidiaGPUDefaults()
	if err != nil {
		t.Fatalf("err=%v", err)
	}
	if out == "" {
		t.Fatalf("expected non-empty output")
	}
	// GPU0: power limit drifted (350 vs 700) -> reset; MIG enabled -> disable; CC ON -> disable.
	// GPU1: power limit already default, MIG already disabled, CC already off;
	// only ECC disabled -> re-enable.
	want := []string{
		"-i 0 -pl 700",
		"-i 0 -mig 0",
		"conf-compute -i 0 -scc 0",
		"-i 1 -e 1",
	}
	if len(calls) != len(want) {
		t.Fatalf("calls=%v want %v", calls, want)
	}
	for i := range want {
		if calls[i] != want[i] {
			t.Fatalf("calls[%d]=%q want %q (all calls: %v)", i, calls[i], want[i], calls)
		}
	}
}

func TestResetNvidiaGPUDefaultsNoOpWhenAlreadyDefault(t *testing.T) {
	var calls []string
	fakeNvidiaSMI(t,
		"0, H100, Enabled, Enabled, Disabled, Disabled, 700.00, 100.00, 700.00, 700.00\n",
		map[int]string{0: "OFF"}, &calls)

	s := &System{}
	out, err := s.ResetNvidiaGPUDefaults()
	if err != nil {
		t.Fatalf("err=%v", err)
	}
	if len(calls) != 0 {
		t.Fatalf("calls=%v want no calls (already at factory defaults)", calls)
	}
	if !strings.Contains(out, "already at factory-default") {
		t.Fatalf("out=%q want no-op message", out)
	}
}