From 14781c7fc6bea0a71ab5a72d750079060aa7a1cc Mon Sep 17 00:00:00 2001
From: Carson Ip <carson.ip@elastic.co>
Date: Wed, 24 Jan 2024 11:41:42 +0000
Subject: [PATCH] Override RUM data stream fields

---
 model/modelprocessor/datastream.go      |  4 ++--
 model/modelprocessor/datastream_test.go | 16 +++++++++-------
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/model/modelprocessor/datastream.go b/model/modelprocessor/datastream.go
index 6ca447ec..b7f14a59 100644
--- a/model/modelprocessor/datastream.go
+++ b/model/modelprocessor/datastream.go
@@ -84,13 +84,13 @@ func (s *SetDataStream) setDataStream(event *modelpb.APMEvent) {
 		event.DataStream.Dataset = errorsDataset
 	case modelpb.LogEventType:
 		event.DataStream.Type = logsType
-		if event.DataStream.Dataset == "" {
+		if event.DataStream.Dataset == "" || isRUMAgentName(event.GetAgent().GetName()) {
 			// Only set dataset if it is not already set in the input event
 			event.DataStream.Dataset = getAppLogsDataset(event)
 		}
 	case modelpb.MetricEventType:
 		event.DataStream.Type = metricsType
-		if event.DataStream.Dataset == "" {
+		if event.DataStream.Dataset == "" || isRUMAgentName(event.GetAgent().GetName()) {
 			// Only set dataset if it is not already set in the input event
 			event.DataStream.Dataset = metricsetDataset(event)
 		}
diff --git a/model/modelprocessor/datastream_test.go b/model/modelprocessor/datastream_test.go
index 8f24e1ef..75013922 100644
--- a/model/modelprocessor/datastream_test.go
+++ b/model/modelprocessor/datastream_test.go
@@ -106,7 +106,7 @@ func TestSetDataStream(t *testing.T) {
 			Service:    &modelpb.Service{Name: "service-name"},
 			DataStream: &modelpb.DataStream{Dataset: "dataset", Namespace: "namespace"},
 		},
-		output: &modelpb.DataStream{Type: "logs", Dataset: "dataset", Namespace: "namespace"},
+		output: &modelpb.DataStream{Type: "logs", Dataset: "apm.app.service_name", Namespace: "custom"},
 	}, {
 		input: &modelpb.APMEvent{
 			Agent:       &modelpb.Agent{Name: "rum-js"},
@@ -150,7 +150,7 @@ func TestSetDataStream(t *testing.T) {
 			},
 			DataStream: &modelpb.DataStream{Dataset: "dataset", Namespace: "namespace"},
 		},
-		output: &modelpb.DataStream{Type: "metrics", Dataset: "dataset", Namespace: "namespace"},
+		output: &modelpb.DataStream{Type: "metrics", Dataset: "apm.app.service_name", Namespace: "custom"},
 	}, {
 		input: &modelpb.APMEvent{
 			Service:     &modelpb.Service{Name: "service-name"},
@@ -193,11 +193,13 @@ func TestSetDataStream(t *testing.T) {
 	}}
 
 	for _, test := range tests {
-		batch := modelpb.Batch{test.input}
-		processor := modelprocessor.SetDataStream{Namespace: "custom"}
-		err := processor.ProcessBatch(context.Background(), &batch)
-		assert.NoError(t, err)
-		assert.Equal(t, test.output, batch[0].DataStream)
+		t.Run("", func(t *testing.T) {
+			batch := modelpb.Batch{test.input}
+			processor := modelprocessor.SetDataStream{Namespace: "custom"}
+			err := processor.ProcessBatch(context.Background(), &batch)
+			assert.NoError(t, err)
+			assert.Equal(t, test.output, batch[0].DataStream)
+		})
 	}
 
 }