Improve PCIe inventory and hardware identity collection

audit/internal/collector/pcie_identity_test.go

@@ -0,0 +1,47 @@
+package collector
+
+import (
+	"bee/audit/internal/schema"
+	"fmt"
+	"testing"
+)
+
+func TestEnrichPCIeWithPCISerialsAddsGPUFallback(t *testing.T) {
+	origDetail := queryPCILSPCIDetail
+	origVPD := readPCIVPDFile
+	t.Cleanup(func() {
+		queryPCILSPCIDetail = origDetail
+		readPCIVPDFile = origVPD
+	})
+
+	queryPCILSPCIDetail = func(bdf string) (string, error) {
+		if bdf != "0000:11:00.0" {
+			t.Fatalf("unexpected bdf: %s", bdf)
+		}
+		return "Serial number: GPU-SN-12345\n", nil
+	}
+	readPCIVPDFile = func(string) ([]byte, error) {
+		return nil, fmt.Errorf("no vpd needed")
+	}
+
+	class := "DisplayController"
+	bdf := "0000:11:00.0"
+	devs := []schema.HardwarePCIeDevice{{
+		DeviceClass: &class,
+		BDF:         &bdf,
+	}}
+
+	out := enrichPCIeWithPCISerials(devs)
+	if out[0].SerialNumber == nil || *out[0].SerialNumber != "GPU-SN-12345" {
+		t.Fatalf("serial=%v want GPU-SN-12345", out[0].SerialNumber)
+	}
+}
+
+func TestShouldProbePCIeSerialSkipsNonGPUOrNIC(t *testing.T) {
+	class := "StorageController"
+	bdf := "0000:19:00.0"
+	dev := schema.HardwarePCIeDevice{DeviceClass: &class, BDF: &bdf}
+	if shouldProbePCIeSerial(dev) {
+		t.Fatal("unexpected probe for storage controller")
+	}
+}