fix bug found with the system-tests

Signed-off-by: Eliott Bouhana <[email protected]>
DataDog · Feb 3, 2025 · 09cab87 · 09cab87
1 parent 373e2fd
commit 09cab87
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 26 deletions.
diff --git a/internal/newtelemetry/client_test.go b/internal/newtelemetry/client_test.go
@@ -117,11 +117,11 @@ func TestClientFlush(t *testing.T) {
 				payload := payloads[0]
 				require.IsType(t, transport.MessageBatch{}, payload)
 				batch := payload.(transport.MessageBatch)
-				require.Len(t, batch.Payload, 2)
-				assert.Equal(t, transport.RequestTypeAppClientConfigurationChange, batch.Payload[0].RequestType)
-				assert.Equal(t, transport.RequestTypeAppExtendedHeartBeat, batch.Payload[1].RequestType)
+				require.Len(t, batch, 2)
+				assert.Equal(t, transport.RequestTypeAppClientConfigurationChange, batch[0].RequestType)
+				assert.Equal(t, transport.RequestTypeAppExtendedHeartBeat, batch[1].RequestType)
 
-				assert.Len(t, batch.Payload[1].Payload.(transport.AppExtendedHeartbeat).Configuration, 0)
+				assert.Len(t, batch[1].Payload.(transport.AppExtendedHeartbeat).Configuration, 0)
 			},
 		},
 		{
@@ -137,12 +137,12 @@ func TestClientFlush(t *testing.T) {
 				payload := payloads[0]
 				require.IsType(t, transport.MessageBatch{}, payload)
 				batch := payload.(transport.MessageBatch)
-				require.Len(t, batch.Payload, 2)
-				assert.Equal(t, transport.RequestTypeAppIntegrationsChange, batch.Payload[0].RequestType)
-				assert.Equal(t, transport.RequestTypeAppExtendedHeartBeat, batch.Payload[1].RequestType)
-				assert.Len(t, batch.Payload[1].Payload.(transport.AppExtendedHeartbeat).Integrations, 1)
-				assert.Equal(t, batch.Payload[1].Payload.(transport.AppExtendedHeartbeat).Integrations[0].Name, "test-integration")
-				assert.Equal(t, batch.Payload[1].Payload.(transport.AppExtendedHeartbeat).Integrations[0].Version, "1.0.0")
+				require.Len(t, batch, 2)
+				assert.Equal(t, transport.RequestTypeAppIntegrationsChange, batch[0].RequestType)
+				assert.Equal(t, transport.RequestTypeAppExtendedHeartBeat, batch[1].RequestType)
+				assert.Len(t, batch[1].Payload.(transport.AppExtendedHeartbeat).Integrations, 1)
+				assert.Equal(t, batch[1].Payload.(transport.AppExtendedHeartbeat).Integrations[0].Name, "test-integration")
+				assert.Equal(t, batch[1].Payload.(transport.AppExtendedHeartbeat).Integrations[0].Version, "1.0.0")
 			},
 		},
 		{
@@ -296,8 +296,8 @@ func TestClientFlush(t *testing.T) {
 				payload := payloads[0]
 				require.IsType(t, transport.MessageBatch{}, payload)
 				batch := payload.(transport.MessageBatch)
-				assert.Len(t, batch.Payload, 2)
-				for _, payload := range batch.Payload {
+				assert.Len(t, batch, 2)
+				for _, payload := range batch {
 					switch p := payload.Payload.(type) {
 					case transport.AppProductChange:
 						assert.Equal(t, transport.RequestTypeAppProductChange, payload.RequestType)
@@ -328,8 +328,8 @@ func TestClientFlush(t *testing.T) {
 				payload := payloads[0]
 				require.IsType(t, transport.MessageBatch{}, payload)
 				batch := payload.(transport.MessageBatch)
-				assert.Len(t, batch.Payload, 3)
-				for _, payload := range batch.Payload {
+				assert.Len(t, batch, 3)
+				for _, payload := range batch {
 					switch p := payload.Payload.(type) {
 					case transport.AppProductChange:
 						assert.Equal(t, transport.RequestTypeAppProductChange, payload.RequestType)
@@ -1097,6 +1097,21 @@ func TestClientEnd2End(t *testing.T) {
 				assert.Equal(t, transport.RequestTypeAppClosing, bodies[1].RequestType)
 			},
 		},
+		{
+			name: "message-batch",
+			when: func(c *client) {
+				c.RegisterAppConfig("key", "value", OriginCode)
+				c.ProductStarted(NamespaceAppSec)
+			},
+			expect: func(t *testing.T, bodies []transport.Body) {
+				require.Len(t, bodies, 1)
+				assert.Equal(t, transport.RequestTypeMessageBatch, bodies[0].RequestType)
+				batch := bodies[0].Payload.(transport.MessageBatch)
+				require.Len(t, batch, 2)
+				assert.Equal(t, transport.RequestTypeAppProductChange, batch[0].RequestType)
+				assert.Equal(t, transport.RequestTypeAppClientConfigurationChange, batch[1].RequestType)
+			},
+		},
 		{
 			name: "fail-agent-endpoint",
 			when: func(c *client) {

diff --git a/internal/newtelemetry/internal/mapper/default.go b/internal/newtelemetry/internal/mapper/default.go
@@ -55,7 +55,7 @@ func (t *messageBatchReducer) Transform(payloads []transport.Payload) ([]transpo
 		}
 	}
 
-	return []transport.Payload{transport.MessageBatch{Payload: messages}}, t
+	return []transport.Payload{transport.MessageBatch(messages)}, t
 }
 
 type heartbeatEnricher struct {

diff --git a/internal/newtelemetry/internal/transport/body.go b/internal/newtelemetry/internal/transport/body.go
@@ -89,26 +89,24 @@ func (b *Body) UnmarshalJSON(bytes []byte) error {
 	}
 
 	if b.RequestType == RequestTypeMessageBatch {
-		var messageBatch struct {
-			Payload []struct {
-				RequestType RequestType     `json:"request_type"`
-				Payload     json.RawMessage `json:"payload"`
-			} `json:"payload"`
+		var messageBatch []struct {
+			RequestType RequestType     `json:"request_type"`
+			Payload     json.RawMessage `json:"payload"`
 		}
 
 		if err = json.Unmarshal(anyMap["payload"], &messageBatch); err != nil {
 			return err
 		}
 
-		batch := make([]Message, len(messageBatch.Payload))
-		for i, message := range messageBatch.Payload {
+		batch := make([]Message, len(messageBatch))
+		for i, message := range messageBatch {
 			payload, err := unmarshalPayload(message.Payload, message.RequestType)
 			if err != nil {
 				return err
 			}
 			batch[i] = Message{RequestType: message.RequestType, Payload: payload}
 		}
-		b.Payload = MessageBatch{Payload: batch}
+		b.Payload = MessageBatch(batch)
 		return nil
 	}
 

diff --git a/internal/newtelemetry/internal/transport/message_batch.go b/internal/newtelemetry/internal/transport/message_batch.go
@@ -5,9 +5,7 @@
 
 package transport
 
-type MessageBatch struct {
-	Payload []Message `json:"payload"`
-}
+type MessageBatch []Message
 
 func (MessageBatch) RequestType() RequestType {
 	return RequestTypeMessageBatch
-Original file line number
+Diff line change
@@ Expand Up @@
     		}
     	}
-    	return []transport.Payload{transport.MessageBatch{Payload: messages}}, t
+    	return []transport.Payload{transport.MessageBatch(messages)}, t
     }
     type heartbeatEnricher struct {
@@ Expand Down @@