More Data added to OHI Inventory

pkg/integrations/legacy/runner.go

@@ -707,9 +707,11 @@ func EmitDataSet(
 	if err != nil {
 		return fmt.Errorf("couldn't determine a unique entity Key: %s", err.Error())
 	}
+	// Custom attributes are from infra agent config
+	customAttr := ctx.Config().CustomAttributes.DataMap()
 
 	if len(dataSet.Inventory) > 0 {
-		inventoryDataSet := BuildInventoryDataSet(elog, dataSet.Inventory, labels, integrationUser, pluginName, entityKey.String())
+		inventoryDataSet := BuildInventoryDataSet(elog, dataSet.Inventory, labels, customAttr, integrationUser, pluginName, pluginVersion, agentIdentifier, entityKey.String())
 		emitter.EmitInventory(inventoryDataSet, entity.NewWithoutID(entityKey))
 	}
 

pkg/integrations/legacy/runner_test.go

@@ -509,7 +509,7 @@ func newFakePlugin(ctx customContext, pluginVersion int) externalPlugin {
 				Context: ctx,
 			},
 			pluginInstance: &PluginV1Instance{
-				Labels: map[string]string{"role": "fileserver", "environment": "development"},
+				Labels: map[string]string{"role": "fileserver", "environment": "development", "agent_role": "overwrite agent role"},
 				plugin: &Plugin{
 					Name:            "new-plugin",
 					ProtocolVersion: pluginVersion,
@@ -573,7 +573,7 @@ func newFakePluginWithEnvVars(pluginVersion int) externalPlugin {
 			Context: ctx,
 		},
 		pluginInstance: &PluginV1Instance{
-			Labels: map[string]string{"role": "fileserver", "environment": "development"},
+			Labels: map[string]string{"role": "fileserver", "environment": "development", "agent_role": "overwrite agent role"},
 			plugin: &Plugin{
 				Name:            "new-plugin",
 				ProtocolVersion: pluginVersion,
@@ -676,22 +676,36 @@ func (rs *RunnerSuite) TestPluginHandleOutputV1(c *C) {
 	c.Assert(err, IsNil)
 
 	c.Assert(rd, NotNil)
-	c.Assert(len(rd.Data), Equals, 4)
-	c.Assert(rd.Data[0].SortKey(), Equals, "first")
-
-	invData := rd.Data[3].(protocol.InventoryData)
-	c.Assert(invData["id"], Equals, "integrationUser")
-	c.Assert(invData["value"], Equals, "test")
 
 	c.Assert(event, NotNil)
 	c.Assert(event["event_type"], Equals, "LoadBalancerSample")
 	c.Assert(event["id"], Equals, "first")
 	c.Assert(event["value"], Equals, "random")
 	c.Assert(event["integrationUser"], Equals, "test")
 
-	for _, labelKey := range labelKeys {
-		if rd.Data[1].SortKey() != labelKey && rd.Data[2].SortKey() != labelKey {
-			c.Errorf("There isn't label '%s'' in the inventory", labelKey)
+	expectedLabelValues := map[string]string{
+		"first":              "fake",
+		"labels/my_group":    "test group",
+		"labels/role":        "fileserver",
+		"labels/environment": "development",
+		"labels/agent_role":  "overwrite agent role",
+		"integrationUser":    "test",
+		"integrationName":    "test",
+		"integrationVersion": "1.0.0",
+	}
+
+	c.Assert(len(rd.Data), Equals, 8)
+
+	for _, item := range rd.Data {
+		if invData, ok := item.(protocol.InventoryData); ok {
+			id, _ := invData["id"].(string)
+			value, _ := invData["value"].(string)
+
+			if expectedValue, exists := expectedLabelValues[id]; exists {
+				c.Assert(value, Equals, expectedValue)
+			} else {
+				c.Fatalf("Expected label: %v not found in Inventory", id)
+			}
 		}
 	}
 }
@@ -726,6 +740,7 @@ func (rs *RunnerSuite) TestPluginHandleOutputEventsV1(c *C) {
 	// labels from pluginInstance
 	c.Assert(event["label.environment"], Equals, "development")
 	c.Assert(event["label.role"], Equals, "fileserver")
+	c.Assert(event["label.agent_role"], Equals, "overwrite agent role")
 
 	// labels from databind
 	c.Assert(event["label.expected"], Equals, "extra label")
@@ -831,7 +846,7 @@ func (rs *RunnerSuite) TestEventsPluginRunV1(c *C) {
 	c.Assert(err, IsNil)
 
 	c.Assert(rd, NotNil)
-	c.Assert(len(rd.Data), Equals, 3)
+	c.Assert(len(rd.Data), Equals, 7)
 	c.Assert(rd.Data[0].SortKey(), Equals, "first")
 
 	c.Assert(event, NotNil)
@@ -1030,7 +1045,7 @@ func (rs *RunnerSuite) TestHandleOutputV1(c *C) {
 	c.Assert(err, IsNil)
 
 	// labels are added as inventory
-	c.Assert(len(rd.Data)-len(labelKeys), Equals, 3)
+	c.Assert(len(rd.Data)-len(labelKeys), Equals, 7)
 
 	firstData := rd.Data[0]
 	inv := firstData.(protocol.InventoryData)
@@ -1545,12 +1560,16 @@ func TestParsePayloadV3(t *testing.T) {
 type fakeEmitter struct {
 	lastEventData map[string]interface{}
 	lastEntityKey string
+	inventory     types.PluginInventoryDataset
 }
 
 func (f *fakeEmitter) EmitInventoryWithPluginId(data types.PluginInventoryDataset, entityKey string, pluginId ids.PluginID) {
+	f.inventory = data
 }
 
-func (f *fakeEmitter) EmitInventory(data types.PluginInventoryDataset, entity entity.Entity) {}
+func (f *fakeEmitter) EmitInventory(data types.PluginInventoryDataset, entity entity.Entity) {
+	f.inventory = data
+}
 
 func (f *fakeEmitter) EmitEvent(eventData map[string]interface{}, entityKey entity.Key) {
 	f.lastEventData = eventData
@@ -1586,6 +1605,29 @@ func TestEmitPayloadV2NoDisplayNameNoEntityName(t *testing.T) {
 	assert.EqualValues(t, "Motorbike", emitter.lastEventData["entityName"])
 	assert.EqualValues(t, "motorbike:street_hawk", emitter.lastEventData["entityKey"])
 
+	expectedLabelValues := map[string]string{
+		"integrationUser":    "testuser",
+		"integrationName":    "test/test",
+		"integrationVersion": "x.y.z",
+		"reportingAgent":     "my-agent-id",
+		"motor":              "", // we are type converting to string so motor does not have a value
+	}
+
+	assert.EqualValues(t, len(emitter.inventory), 5)
+
+	for _, item := range emitter.inventory {
+		if invData, ok := item.(protocol.InventoryData); ok {
+			id, _ := invData["id"].(string)
+			value, _ := invData["value"].(string)
+
+			if expectedValue, exists := expectedLabelValues[id]; exists {
+				assert.EqualValues(t, value, expectedValue)
+			} else {
+				assert.Fail(t, "Expected label not found in Inventory", "Label: %s not found in inventory", id)
+			}
+		}
+	}
+
 	// Local entity, no displayName, no entityName
 	assert.NoError(t, EmitDataSet(ctx, &emitter, "test/test", "x.y.z", "testuser", rd.DataSets[2], extraAnnotations, labels, entityRewrite, version))
 	_, ok := emitter.lastEventData["displayName"]
@@ -1594,6 +1636,27 @@ func TestEmitPayloadV2NoDisplayNameNoEntityName(t *testing.T) {
 	assert.False(t, ok)
 	// but entityKey is the agent key
 	assert.EqualValues(t, "my-agent-id", emitter.lastEventData["entityKey"])
+
+	// Check inventory data
+	assert.EqualValues(t, len(emitter.inventory), 5)
+
+	for _, item := range emitter.inventory {
+		if invData, ok := item.(protocol.InventoryData); ok {
+			id, _ := invData["id"].(string)
+			value, _ := invData["value"].(string)
+
+			if expectedValue, exists := expectedLabelValues[id]; exists {
+				assert.EqualValues(t, value, expectedValue)
+			} else {
+				assert.Fail(t, "Expected label not found in Inventory", "Label: %s not found in inventory", id)
+			}
+		}
+	}
+}
+
+func createMockConfigWithDataMap(attrs map[string]interface{}) *config.Config {
+	customAttrs := config.CustomAttributeMap(attrs)
+	return &config.Config{CustomAttributes: customAttrs} //nolint:all
 }
 
 func TestEmitDataSet_OnAddHostnameDecoratesWithHostname(t *testing.T) {
@@ -1626,6 +1689,7 @@ func TestEmitDataSet_OnAddHostnameDecoratesWithHostname(t *testing.T) {
 	ctx.On("EntityKey").Return(agentIdentifier)
 	ctx.On("HostnameResolver").Return(newFixedHostnameResolver(hn, "short"))
 	ctx.On("IDLookup").Return(newFixedIDLookup())
+	ctx.On("Config").Return(createMockConfigWithDataMap(make(map[string]interface{})))
 
 	em := &fakeEmitter{}
 	extraAnnotations := map[string]string{}
@@ -1677,6 +1741,7 @@ func TestEmitDataSet_EntityNameLocalhostIsNotReplacedWithHostnameV2(t *testing.T
 	ctx.On("EntityKey").Return(agID)
 	ctx.On("HostnameResolver").Return(newFixedHostnameResolver("foo.bar", "short"))
 	ctx.On("IDLookup").Return(newFixedIDLookup())
+	ctx.On("Config").Return(createMockConfigWithDataMap(make(map[string]interface{})))
 
 	em := &fakeEmitter{}
 	extraAnnotations := map[string]string{}
@@ -1724,6 +1789,7 @@ func TestEmitDataSet_EntityNameLocalhostIsReplacedWithHostnameV3(t *testing.T) {
 	ctx.On("EntityKey").Return(agID)
 	ctx.On("HostnameResolver").Return(newFixedHostnameResolver("foo.bar", "short"))
 	ctx.On("IDLookup").Return(newFixedIDLookup())
+	ctx.On("Config").Return(createMockConfigWithDataMap(make(map[string]interface{})))
 
 	em := &fakeEmitter{}
 	extraAnnotations := map[string]string{}
@@ -1772,6 +1838,7 @@ func TestEmitDataSet_MetricHostnameIsReplacedIfLocalhostV3(t *testing.T) {
 	ctx.On("EntityKey").Return(agID)
 	ctx.On("HostnameResolver").Return(newFixedHostnameResolver("foo.bar", "short"))
 	ctx.On("IDLookup").Return(newFixedIDLookup())
+	ctx.On("Config").Return(createMockConfigWithDataMap(make(map[string]interface{})))
 
 	em := &fakeEmitter{}
 	extraAnnotations := map[string]string{}
@@ -1821,6 +1888,7 @@ func TestEmitDataSet_ReportingFieldsAreReplacedIfLocalhostV3(t *testing.T) {
 	ctx.On("EntityKey").Return(agID)
 	ctx.On("HostnameResolver").Return(newFixedHostnameResolver("foo.bar", "short"))
 	ctx.On("IDLookup").Return(newFixedIDLookup())
+	ctx.On("Config").Return(createMockConfigWithDataMap(make(map[string]interface{})))
 
 	em := &fakeEmitter{}
 	extraAnnotations := map[string]string{}
@@ -1869,6 +1937,7 @@ func TestEmitDataSet_LogsEntityViolationsOncePerEntity(t *testing.T) {
 	ctx.On("EntityKey").Return(agID)
 	ctx.On("HostnameResolver").Return(newFixedHostnameResolver("foo.bar", "short"))
 	ctx.On("IDLookup").Return(newFixedIDLookup())
+	ctx.On("Config").Return(createMockConfigWithDataMap(make(map[string]interface{})))
 
 	em := &fakeEmitter{}
 	extraAnnotations := map[string]string{}
@@ -1906,6 +1975,7 @@ func TestEmitDataSet_DoNotOverrideExistingMetrics(t *testing.T) {
 	ctx.On("EntityKey").Return("agent-id")
 	ctx.On("HostnameResolver").Return(newFixedHostnameResolver("long", "short"))
 	ctx.On("IDLookup").Return(newFixedIDLookup())
+	ctx.On("Config").Return(createMockConfigWithDataMap(make(map[string]interface{})))
 	em := &fakeEmitter{}
 	extraAnnotations := map[string]string{
 		"cluster_name":       "K8sDiscoveredCluster",
@@ -2207,8 +2277,7 @@ func TestLogFields(t *testing.T) {
 		"TEMP_DIR":   "a/path",
 	})
 	assert.Equal(t, fields["labels"], map[string]string{
-		"role":        "fileserver",
-		"environment": "development",
+		"agent_role": "overwrite agent role", "environment": "development", "role": "fileserver",
 	})
 
 	if runtime.GOOS == "windows" {

pkg/integrations/legacy/utils.go

@@ -4,6 +4,7 @@ package legacy
 
 import (
 	"fmt"
+
 	"github.com/newrelic/infrastructure-agent/internal/agent/types"
 
 	event2 "github.com/newrelic/infrastructure-agent/pkg/event"
@@ -12,12 +13,22 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
+const (
+	integrationUserID    = "integrationUser"
+	integrationNameID    = "integrationName"
+	integrationVersionID = "integrationVersion"
+	reportingAgentID     = "reportingAgent"
+)
+
 func BuildInventoryDataSet(
 	entryLog log.Entry,
 	inventoryData map[string]protocol.InventoryData,
 	labels map[string]string,
+	customAttr map[string]string,
 	integrationUser string,
 	pluginName string,
+	pluginVersion string,
+	reportingAgent string,
 	entityKey string) types.PluginInventoryDataset {
 	var inventoryDataSet types.PluginInventoryDataset
 
@@ -31,20 +42,39 @@ func BuildInventoryDataSet(
 		}
 	}
 
-	for key, value := range labels {
+	addLabeledData := func(id string, value string) {
 		inventoryDataSet = append(inventoryDataSet, protocol.InventoryData{
-			"id":        fmt.Sprintf("labels/%s", key),
+			"id":        id,
 			"value":     value,
 			"entityKey": entityKey,
 		})
 	}
 
+	for key, value := range customAttr {
+		// Do not set in the case of duplicate key
+		if _, exists := labels[key]; !exists {
+			addLabeledData(fmt.Sprintf("labels/%s", key), value)
+		}
+	}
+
+	for key, value := range labels {
+		addLabeledData(fmt.Sprintf("labels/%s", key), value)
+	}
+
 	if integrationUser != "" {
-		inventoryDataSet = append(inventoryDataSet, protocol.InventoryData{
-			"id":        "integrationUser",
-			"value":     integrationUser,
-			"entityKey": entityKey,
-		})
+		addLabeledData(integrationUserID, integrationUser)
+	}
+
+	if pluginName != "" {
+		addLabeledData(integrationNameID, pluginName)
+	}
+
+	if pluginVersion != "" {
+		addLabeledData(integrationVersionID, pluginVersion)
+	}
+
+	if reportingAgent != "" {
+		addLabeledData(reportingAgentID, reportingAgent)
 	}
 
 	return inventoryDataSet