fix: preserve message context items when SetMessage is called

src/KafkaFlow/MessageContext.cs

@@ -11,13 +11,25 @@ public MessageContext(
         IConsumerContext consumer,
         IProducerContext producer,
         IReadOnlyCollection<string> brokers)
+        : this(message, headers, dependencyResolver, consumer, producer, brokers, new Dictionary<string, object>())
+    {
+    }
+
+    private MessageContext(
+        Message message,
+        IMessageHeaders headers,
+        IDependencyResolver dependencyResolver,
+        IConsumerContext consumer,
+        IProducerContext producer,
+        IReadOnlyCollection<string> brokers,
+        IDictionary<string, object> items)
     {
         this.Message = message;
         this.DependencyResolver = dependencyResolver;
         this.Headers = headers ?? new MessageHeaders();
         this.ConsumerContext = consumer;
         this.ProducerContext = producer;
-        this.Items = new Dictionary<string, object>();
+        this.Items = items;
         this.Brokers = brokers;
     }
 
@@ -41,5 +53,6 @@ public MessageContext(
         this.DependencyResolver,
         this.ConsumerContext,
         this.ProducerContext,
-        this.Brokers);
+        this.Brokers,
+        this.Items);
 }

tests/KafkaFlow.UnitTests/MessageContextTests.cs

@@ -0,0 +1,39 @@
+using System.Collections.Generic;
+using KafkaFlow.Consumers;
+using KafkaFlow.Producers;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using Moq;
+
+namespace KafkaFlow.UnitTests;
+
+[TestClass]
+public class MessageContextTests
+{
+    [TestMethod]
+    public void SetMessage_ShouldSetMessageCorrectly()
+    {
+        // Arrange
+        var messageContext = new MessageContext(
+            new Message("key", "value"),
+            Mock.Of<IMessageHeaders>(),
+            Mock.Of<IDependencyResolver>(),
+            Mock.Of<IConsumerContext>(),
+            Mock.Of<IProducerContext>(),
+            Mock.Of<IReadOnlyCollection<string>>()
+        );
+
+
+        // Act
+        var changedMessage = messageContext.SetMessage("changed-key", "changed-value");
+
+        // Assert
+        Assert.AreEqual("changed-key", changedMessage.Message.Key);
+        Assert.AreEqual("changed-value", changedMessage.Message.Value);
+        Assert.AreSame(messageContext.ConsumerContext, changedMessage.ConsumerContext);
+        Assert.AreSame(messageContext.DependencyResolver, changedMessage.DependencyResolver);
+        Assert.AreSame(messageContext.Headers, changedMessage.Headers);
+        Assert.AreSame(messageContext.ProducerContext, changedMessage.ProducerContext);
+        Assert.AreSame(messageContext.Brokers, changedMessage.Brokers);
+        Assert.AreSame(messageContext.Items, changedMessage.Items);
+    }
+}