bee/audit/internal/collector/pcie_nvlink_bridge_test.go

package collector

import (
	"bee/audit/internal/schema"
	"testing"
)

func TestParseNVIDIATopologyMatrix(t *testing.T) {
	t.Parallel()

	// Real-world B200 HGX output: 8 GPUs, all pairs connected via NV18.
	input := `	GPU0	GPU1	GPU2	GPU3	GPU4	GPU5	GPU6	GPU7	NIC0	NIC1
GPU0	 X 	NV18	NV18	NV18	NV18	NV18	NV18	NV18	NODE	NODE
GPU1	NV18	 X 	NV18	NV18	NV18	NV18	NV18	NV18	NODE	NODE
GPU2	NV18	NV18	 X 	NV18	NV18	NV18	NV18	NV18	NODE	NODE
GPU3	NV18	NV18	NV18	 X 	NV18	NV18	NV18	NV18	NODE	NODE
GPU4	NV18	NV18	NV18	NV18	 X 	NV18	NV18	NV18	SYS	SYS
GPU5	NV18	NV18	NV18	NV18	NV18	 X 	NV18	NV18	SYS	SYS
GPU6	NV18	NV18	NV18	NV18	NV18	NV18	 X 	NV18	SYS	SYS
GPU7	NV18	NV18	NV18	NV18	NV18	NV18	NV18	 X 	SYS	SYS
NIC0	NODE	NODE	NODE	NODE	SYS	SYS	SYS	SYS	 X 	PIX
`
	got := parseNVIDIATopologyMatrix(input)

	if got.GPUCount != 8 {
		t.Fatalf("GPUCount=%d want 8", got.GPUCount)
	}
	if !got.AllActive {
		t.Fatalf("AllActive=false want true")
	}
	if got.MinNVLinks != 18 {
		t.Fatalf("MinNVLinks=%d want 18", got.MinNVLinks)
	}
}

func TestParseNVIDIATopologyMatrixPartialDegradation(t *testing.T) {
	t.Parallel()

	// GPU1-GPU3 pair shows NV12 (reduced) instead of NV18.
	input := `	GPU0	GPU1	GPU2	GPU3
GPU0	 X 	NV18	NV18	NV18
GPU1	NV18	 X 	NV18	NV12
GPU2	NV18	NV18	 X 	NV18
GPU3	NV18	NV12	NV18	 X
`
	got := parseNVIDIATopologyMatrix(input)

	if got.MinNVLinks != 12 {
		t.Fatalf("MinNVLinks=%d want 12", got.MinNVLinks)
	}
	if !got.AllActive {
		t.Fatalf("AllActive=false want true (12 links is still active)")
	}
}

func TestParseNVIDIATopologyMatrixDisconnected(t *testing.T) {
	t.Parallel()

	// GPU0-GPU1 pair fully disconnected (NV0).
	input := `	GPU0	GPU1
GPU0	 X 	NV0
GPU1	NV0	 X
`
	got := parseNVIDIATopologyMatrix(input)

	if got.AllActive {
		t.Fatalf("AllActive=true want false (NV0 means no links)")
	}
	if got.MinNVLinks != 0 {
		t.Fatalf("MinNVLinks=%d want 0", got.MinNVLinks)
	}
}

func TestParseNVIDIATopologyMatrixANSIUnderlinedHeader(t *testing.T) {
	t.Parallel()

	// nvidia-smi underlines the header row with ANSI escapes even when stdout
	// is not a TTY: the header line starts with ESC[4m, not "GPU0".
	input := "\x1b[4m\tGPU0\tGPU1\tNIC0\tCPU Affinity\x1b[0m\n" +
		"GPU0\t X \tNV18\tNODE\t0-31,64-95\n" +
		"GPU1\tNV18\t X \tNODE\t0-31,64-95\n" +
		"NIC0\tNODE\tNODE\t X \n"

	got := parseNVIDIATopologyMatrix(input)
	if got.GPUCount != 2 {
		t.Fatalf("GPUCount=%d want 2", got.GPUCount)
	}
	if !got.AllActive || got.MinNVLinks != 18 {
		t.Fatalf("AllActive=%v MinNVLinks=%d want true/18", got.AllActive, got.MinNVLinks)
	}
}

func TestParseNVIDIATopologyMatrixEmpty(t *testing.T) {
	t.Parallel()

	got := parseNVIDIATopologyMatrix("no gpus here")
	if got.GPUCount != 0 {
		t.Fatalf("GPUCount=%d want 0", got.GPUCount)
	}
}

func TestParseCrossNUMAPeersDetectsSYS(t *testing.T) {
	t.Parallel()

	// 4-GPU box, two NVLink-bridged pairs (GPU0-GPU1, GPU2-GPU3); the pairs
	// themselves only reach each other via SYS (cross-NUMA PCIe hop) — the
	// exact topology of a server using pairwise NVLink bridge cards instead
	// of a switched NVLink fabric.
	input := `	GPU0	GPU1	GPU2	GPU3
GPU0	 X 	NV4	SYS	SYS
GPU1	NV4	 X 	SYS	SYS
GPU2	SYS	SYS	 X 	NV4
GPU3	SYS	SYS	NV4	 X
`
	peers := parseCrossNUMAPeers(input)

	if len(peers[0]) != 2 || peers[0][0] != 2 || peers[0][1] != 3 {
		t.Fatalf("peers[0]=%v want [2 3]", peers[0])
	}
	if len(peers[2]) != 2 {
		t.Fatalf("peers[2]=%v want 2 entries", peers[2])
	}
}

func TestParseCrossNUMAPeersNoSYS(t *testing.T) {
	t.Parallel()

	// Full NVSwitch fabric: every GPU pair connects via NVLink, no SYS hops.
	input := `	GPU0	GPU1
GPU0	 X 	NV18
GPU1	NV18	 X
`
	if peers := parseCrossNUMAPeers(input); peers != nil {
		t.Fatalf("peers=%v want nil (no SYS pairs)", peers)
	}
}

func TestApplyPCIeLinkSpeedWarningNVLinkBridgeEscalates(t *testing.T) {
	t.Parallel()

	bridgeClass := "NVLinkBridge"
	linkSpeed := "Gen3"
	maxLinkSpeed := "Gen4"
	dev := schema.HardwarePCIeDevice{}
	dev.DeviceClass = &bridgeClass
	dev.LinkSpeed = &linkSpeed
	dev.MaxLinkSpeed = &maxLinkSpeed
	s := statusOK
	dev.Status = &s

	applyPCIeLinkSpeedWarning(&dev)

	if dev.Status == nil || *dev.Status != statusCritical {
		t.Fatalf("status=%v want Critical for NVLink bridge degradation", dev.Status)
	}
	if dev.ErrorDescription == nil {
		t.Fatal("ErrorDescription nil, want degradation message")
	}
}

func TestApplyPCIeLinkSpeedWarningRegularCardIsWarning(t *testing.T) {
	t.Parallel()

	regularClass := "NetworkController"
	linkSpeed := "Gen3"
	maxLinkSpeed := "Gen4"
	dev := schema.HardwarePCIeDevice{}
	dev.DeviceClass = &regularClass
	dev.LinkSpeed = &linkSpeed
	dev.MaxLinkSpeed = &maxLinkSpeed
	s := statusOK
	dev.Status = &s

	applyPCIeLinkSpeedWarning(&dev)

	if dev.Status == nil || *dev.Status != statusWarning {
		t.Fatalf("status=%v want Warning for regular card degradation", dev.Status)
	}
}