Add test for the collector

comp/core/workloadmeta/collectors/internal/nvml/nvml.go

@@ -23,6 +23,7 @@ import (
 const (
 	collectorID   = "nvml"
 	componentName = "workloadmeta-nvml"
+	nvidiaVendor  = "nvidia"
 )
 
 type collector struct {
@@ -96,7 +97,7 @@ func (c *collector) Pull(_ context.Context) error {
 			EntityMeta: workloadmeta.EntityMeta{
 				Name: name,
 			},
-			Vendor: "nvidia",
+			Vendor: nvidiaVendor,
 			Device: name,
 		}
 

comp/core/workloadmeta/collectors/internal/nvml/nvml_test.go

@@ -0,0 +1,48 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2024-present Datadog, Inc.
+
+//go:build linux
+
+package nvml
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	workloadmeta "github.com/DataDog/datadog-agent/comp/core/workloadmeta/def"
+	"github.com/DataDog/datadog-agent/pkg/gpu/testutil"
+)
+
+func TestPull(t *testing.T) {
+	wmetaMock := testutil.GetWorkloadMetaMock(t)
+	nvmlMock := testutil.GetBasicNvmlMock()
+
+	c := &collector{
+		id:      collectorID,
+		catalog: workloadmeta.NodeAgent,
+		store:   wmetaMock,
+		nvmlLib: nvmlMock,
+	}
+
+	c.Pull(context.Background())
+
+	gpus := wmetaMock.ListGPUs()
+	require.Equal(t, len(testutil.GPUUUIDs), len(gpus))
+
+	foundIDs := make(map[string]bool)
+	for _, gpu := range gpus {
+		foundIDs[gpu.ID] = true
+
+		require.Equal(t, nvidiaVendor, gpu.Vendor)
+		require.Equal(t, testutil.DefaultGPUName, gpu.Name)
+		require.Equal(t, testutil.DefaultGPUName, gpu.Device)
+	}
+
+	for _, uuid := range testutil.GPUUUIDs {
+		require.True(t, foundIDs[uuid], "GPU with UUID %s not found", uuid)
+	}
+}

pkg/gpu/testutil/mocks.go

@@ -45,6 +45,9 @@ var GPUCores = []int{DefaultGpuCores, 20, 30, 40, 50, 60, 70}
 // DefaultGpuUUID is the UUID for the default device returned by the mock
 var DefaultGpuUUID = GPUUUIDs[0]
 
+// DefaultGPUName is the name for the default device returned by the mock
+var DefaultGPUName = "Tesla T4"
+
 // GetDeviceMock returns a mock of the nvml.Device with the given UUID.
 func GetDeviceMock(deviceIdx int) *nvmlmock.Device {
 	return &nvmlmock.Device{
@@ -57,6 +60,9 @@ func GetDeviceMock(deviceIdx int) *nvmlmock.Device {
 		GetUUIDFunc: func() (string, nvml.Return) {
 			return GPUUUIDs[deviceIdx], nvml.SUCCESS
 		},
+		GetNameFunc: func() (string, nvml.Return) {
+			return DefaultGPUName, nvml.SUCCESS
+		},
 	}
 }