From 4f79c43ca3e2a06d30114211c1ab29ceaefa003d Mon Sep 17 00:00:00 2001
From: Sascha Kiefer <esskar@gmail.com>
Date: Tue, 27 Feb 2024 09:24:13 +0100
Subject: [PATCH] fix: preserve message context items when SetMessage is called

---
 src/KafkaFlow/MessageContext.cs               | 17 +++++++-
 .../MessageContextTests.cs                    | 39 +++++++++++++++++++
 2 files changed, 54 insertions(+), 2 deletions(-)
 create mode 100644 tests/KafkaFlow.UnitTests/MessageContextTests.cs

diff --git a/src/KafkaFlow/MessageContext.cs b/src/KafkaFlow/MessageContext.cs
index bcb406d70..da0e66d77 100644
--- a/src/KafkaFlow/MessageContext.cs
+++ b/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);
 }
diff --git a/tests/KafkaFlow.UnitTests/MessageContextTests.cs b/tests/KafkaFlow.UnitTests/MessageContextTests.cs
new file mode 100644
index 000000000..442b8ad34
--- /dev/null
+++ b/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);
+    }
+}
\ No newline at end of file