RabbitMqContextManager is resilient to missing headers

Nybus-project · Oct 12, 2015 · 29c9201 · 29c9201
1 parent 21e59c1
commit 29c9201
Show file tree

Hide file tree

Showing 2 changed files with 89 additions and 2 deletions.
diff --git a/src/MassTransit/Configuration/IContextManager.cs b/src/MassTransit/Configuration/IContextManager.cs
@@ -58,7 +58,14 @@ public class RabbitMqContextManager : IContextManager
 
         private Guid ExtractCorrelationId(IMessageHeaders messageHeaders)
         {
-            return Guid.Parse(messageHeaders[CorrelationIdKey]);
+            var header = messageHeaders[CorrelationIdKey];
+
+            if (header == null)
+            {
+                return Guid.NewGuid();
+            }
+
+            return Guid.Parse(header);
         }
 
         private void PersistCorrelationId(ISendContext sendContext, Guid correlationId)
@@ -68,7 +75,14 @@ private void PersistCorrelationId(ISendContext sendContext, Guid correlationId)
 
         private DateTimeOffset ExtractMessageSentTime(IMessageHeaders messageHeaders)
         {
-            return DateTimeOffset.Parse(messageHeaders[MessageSentKey]);
+            var header = messageHeaders[MessageSentKey];
+
+            if (header == null)
+            {
+                return Clock.Default.Now;
+            }
+
+            return DateTimeOffset.Parse(header);
         }
 
         private void PersistMessageSentTime(ISendContext sendContext, DateTimeOffset sentTime)

diff --git a/tests/Tests.MassTransit/Configuration/RabbitMqContextManagerTests.cs b/tests/Tests.MassTransit/Configuration/RabbitMqContextManagerTests.cs
@@ -4,6 +4,7 @@
 using NUnit.Framework;
 using Nybus;
 using Nybus.Configuration;
+using Nybus.Utils;
 using Ploeh.AutoFixture;
 
 namespace Tests.Configuration
@@ -22,6 +23,8 @@ public class RabbitMqContextManagerTests
 
         private DateTimeOffset now;
 
+        private Mock<IClock> mockClock;
+
         [SetUp]
         public void Initialize()
         {
@@ -36,6 +39,11 @@ public void Initialize()
             mockCommandPublishContext = new Mock<IPublishContext<TestCommand>>();
 
             now = fixture.Create<DateTimeOffset>();
+
+            mockClock = new Mock<IClock>();
+            mockClock.SetupGet(p => p.Now).Returns(now);
+
+            Clock.Default = mockClock.Object;
         }
 
         private RabbitMqContextManager CreateSystemUnderTest()
@@ -79,6 +87,38 @@ public void CreateEventMessage_retrieves_Event_from_context()
             Assert.That(message.Event, Is.SameAs(body));
         }
 
+        [Test]
+        public void CreateEventMessage_returns_Now_if_SentOn_is_missing()
+        {
+            var sut = CreateSystemUnderTest();
+
+            TestEvent body = fixture.Create<TestEvent>();
+
+            Guid correlationId = fixture.Create<Guid>();
+
+            mockEventConsumeContext.SetupGet(p => p.Message).Returns(body);
+            mockEventConsumeContext.Setup(p => p.Headers[RabbitMqContextManager.CorrelationIdKey]).Returns(correlationId.ToString());
+
+            var message = sut.CreateEventMessage(mockEventConsumeContext.Object);
+
+            Assert.That(message.SentOn, Is.EqualTo(now));
+        }
+
+        [Test]
+        public void CreateEventMessage_returns_NewGuid_if_CorrelationId_is_missing()
+        {
+            var sut = CreateSystemUnderTest();
+
+            TestEvent body = fixture.Create<TestEvent>();
+
+            mockEventConsumeContext.SetupGet(p => p.Message).Returns(body);
+            mockEventConsumeContext.Setup(p => p.Headers[RabbitMqContextManager.MessageSentKey]).Returns(now.ToString("O"));
+
+            var message = sut.CreateEventMessage(mockEventConsumeContext.Object);
+
+            Assert.That(message.CorrelationId, Is.Not.EqualTo(Guid.Empty));
+        }
+
         [Test]
         public void SetEventMessageHeaders_sets_correlationId_to_context()
         {
@@ -130,6 +170,39 @@ public void CreateCommandMessage_retrieves_Command_from_context()
             Assert.That(message.SentOn, Is.EqualTo(now));
         }
 
+        [Test]
+        public void CreateCommandMessage_returns_Now_if_SentOn_is_missing()
+        {
+            var sut = CreateSystemUnderTest();
+
+            var body = fixture.Create<TestCommand>();
+
+            var correlationId = fixture.Create<Guid>();
+
+            mockCommandConsumeContext.SetupGet(p => p.Message).Returns(body);
+            mockCommandConsumeContext.Setup(p => p.Headers[RabbitMqContextManager.CorrelationIdKey]).Returns(correlationId.ToString());
+
+            var message = sut.CreateCommandMessage(mockCommandConsumeContext.Object);
+
+            Assert.That(message.SentOn, Is.EqualTo(now));
+        }
+
+        [Test]
+        public void CreateCommandMessage_returns_NewGuid_if_CorrelationId_is_missing()
+        {
+            var sut = CreateSystemUnderTest();
+
+            var body = fixture.Create<TestCommand>();
+
+            mockCommandConsumeContext.SetupGet(p => p.Message).Returns(body);
+            mockCommandConsumeContext.Setup(p => p.Headers[RabbitMqContextManager.MessageSentKey]).Returns(now.ToString("O"));
+
+            var message = sut.CreateCommandMessage(mockCommandConsumeContext.Object);
+
+            Assert.That(message.CorrelationId, Is.Not.EqualTo(Guid.Empty));
+        }
+
+
         [Test]
         public void SetCommandMessageHeaders_sets_correlationId_to_context()
         {