package main

import "testing"

func TestParseGPUMemSingleDiscreteGPU(t *testing.T) {
	g, ok := parseGPUMem("81920, 1024\n")
	if !ok || g.Total != 81920<<20 || g.Used != 1024<<20 || g.GPUs != 1 || g.MaxOne != 81920<<20 {
		t.Fatalf("%+v %v", g, ok)
	}
	if g.avail() != (81920-1024)<<20 {
		t.Fatal(g.avail())
	}
}

func TestParseGPUMemSumsSeveralGPUs(t *testing.T) {
	g, ok := parseGPUMem("81920, 80000\n81920, 0\n40960, 100\n")
	if !ok || g.GPUs != 3 || g.Total != (81920+81920+40960)<<20 || g.MaxOne != 81920<<20 {
		t.Fatalf("%+v", g)
	}
	if g.avail() != (1920+81920+40860)<<20 {
		t.Fatalf("free memory is total - used: %d", g.avail())
	}
}

func TestUnifiedMemoryMachinesReportNotAvailableAndAreNotTreatedAsVRAM(t *testing.T) {
	// DGX Spark / GB10: nvidia-smi prints [N/A] for memory, so the budget must stay the system RAM
	if _, ok := parseGPUMem("[N/A], [N/A]\n"); ok {
		t.Fatal("N/A must not count as GPU memory")
	}
	if _, ok := parseGPUMem("[Not Supported], [Not Supported]\n"); ok {
		t.Fatal("Not Supported must not count as GPU memory")
	}
	if _, ok := parseGPUMem(""); ok {
		t.Fatal("no output means no GPU")
	}
}

func TestMixedLinesKeepOnlyTheNumericGPUs(t *testing.T) {
	g, ok := parseGPUMem("[N/A], [N/A]\n24576, 2048\n")
	if !ok || g.GPUs != 1 || g.Total != 24576<<20 {
		t.Fatalf("%+v", g)
	}
}

func TestAvailNeverUnderflows(t *testing.T) {
	g, _ := parseGPUMem("1000, 1500\n")
	if g.avail() != 0 {
		t.Fatalf("used > total must give 0, got %d", g.avail())
	}
}