From 1310be6ace26fce0e7e184238edb50e4c451a80c Mon Sep 17 00:00:00 2001
From: TravisNickels <87037242+TravisNickels@users.noreply.github.com>
Date: Tue, 22 Feb 2022 20:17:29 -0800
Subject: [PATCH] Fixed how the DeliverAt header is assigned to a
 DelayedDeliveryConstraint

---
 .../Outgoing/AttachSenderRelatedInfoOnMessageTests.cs      | 3 ++-
 .../Outgoing/AttachSenderRelatedInfoOnMessageBehavior.cs   | 7 ++++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/NServiceBus.Core.Tests/Pipeline/Outgoing/AttachSenderRelatedInfoOnMessageTests.cs b/src/NServiceBus.Core.Tests/Pipeline/Outgoing/AttachSenderRelatedInfoOnMessageTests.cs
index ac00b067891..7e9f9cdc677 100644
--- a/src/NServiceBus.Core.Tests/Pipeline/Outgoing/AttachSenderRelatedInfoOnMessageTests.cs
+++ b/src/NServiceBus.Core.Tests/Pipeline/Outgoing/AttachSenderRelatedInfoOnMessageTests.cs
@@ -3,6 +3,7 @@
     using System;
     using System.Collections.Generic;
     using System.Threading.Tasks;
+    using DeliveryConstraints;
     using Extensibility;
     using NServiceBus.Routing;
     using NUnit.Framework;
@@ -103,7 +104,7 @@ static async Task<OutgoingMessage> InvokeBehavior(Dictionary<string, string> hea
 
             if (options != null)
             {
-                stash.Set(options);
+                stash.AddDeliveryConstraint(options.DelayedDeliveryConstraint);
             }
 
             await new AttachSenderRelatedInfoOnMessageBehavior()
diff --git a/src/NServiceBus.Core/Pipeline/Outgoing/AttachSenderRelatedInfoOnMessageBehavior.cs b/src/NServiceBus.Core/Pipeline/Outgoing/AttachSenderRelatedInfoOnMessageBehavior.cs
index 20cb07497d7..7ce473357bd 100644
--- a/src/NServiceBus.Core/Pipeline/Outgoing/AttachSenderRelatedInfoOnMessageBehavior.cs
+++ b/src/NServiceBus.Core/Pipeline/Outgoing/AttachSenderRelatedInfoOnMessageBehavior.cs
@@ -3,6 +3,7 @@ namespace NServiceBus
     using System;
     using System.Threading.Tasks;
     using DelayedDelivery;
+    using NServiceBus.DeliveryConstraints;
     using Pipeline;
 
     class AttachSenderRelatedInfoOnMessageBehavior : IBehavior<IRoutingContext, IRoutingContext>
@@ -27,14 +28,14 @@ public Task Invoke(IRoutingContext context, Func<IRoutingContext, Task> next)
 
             if (!message.Headers.ContainsKey(Headers.DeliverAt))
             {
-                if (context.Extensions.TryGet<SendOptions>(out var options))
+                if (context.Extensions.TryGetDeliveryConstraint<DelayedDeliveryConstraint>(out var delayedDeliveryConstraint))
                 {
-                    if (options.DelayedDeliveryConstraint is DelayDeliveryWith delayDeliveryWith)
+                    if (delayedDeliveryConstraint is DelayDeliveryWith delayDeliveryWith)
                     {
                         var timeDelay = delayDeliveryWith.Delay;
                         message.Headers[Headers.DeliverAt] = DateTimeExtensions.ToWireFormattedString(utcNow.Add(timeDelay));
                     }
-                    else if (options.DelayedDeliveryConstraint is DoNotDeliverBefore doNotDeliverBefore)
+                    else if (delayedDeliveryConstraint is DoNotDeliverBefore doNotDeliverBefore)
                     {
                         var deliverAt = doNotDeliverBefore.At;
                         message.Headers[Headers.DeliverAt] = DateTimeExtensions.ToWireFormattedString(deliverAt);