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")
	}
}