From 4a4eab82ef0abb378c70a758a49d9f8884be599c Mon Sep 17 00:00:00 2001
From: Yi Liu <yiliu6@microsoft.com>
Date: Thu, 16 Jul 2020 15:07:20 +0800
Subject: [PATCH] enable scheduled enqueue message in servicebus binder

---
 .../azure/spring/integration/core/AzureHeaders.java |  2 ++
 .../converter/ServiceBusMessageConverter.java       |  8 ++++++++
 .../servicebus/ServiceBusMessageConverterTest.java  | 13 +++++++++++++
 3 files changed, 23 insertions(+)

diff --git a/spring-integration-azure/spring-integration-azure-core/src/main/java/com/microsoft/azure/spring/integration/core/AzureHeaders.java b/spring-integration-azure/spring-integration-azure-core/src/main/java/com/microsoft/azure/spring/integration/core/AzureHeaders.java
index 925b36bd..b06d38b8 100644
--- a/spring-integration-azure/spring-integration-azure-core/src/main/java/com/microsoft/azure/spring/integration/core/AzureHeaders.java
+++ b/spring-integration-azure/spring-integration-azure-core/src/main/java/com/microsoft/azure/spring/integration/core/AzureHeaders.java
@@ -22,6 +22,8 @@ public class AzureHeaders {
 
     public static final String NAME = PREFIX + "name";
 
+    public static final String SCHEDULED_ENQUEUE_MESSAGE = "x-delay";
+
     /**
      * The {@value CHECKPOINTER} header for checkpoint the specific message.
      */
diff --git a/spring-integration-azure/spring-integration-servicebus/src/main/java/com/microsoft/azure/spring/integration/servicebus/converter/ServiceBusMessageConverter.java b/spring-integration-azure/spring-integration-servicebus/src/main/java/com/microsoft/azure/spring/integration/servicebus/converter/ServiceBusMessageConverter.java
index 945db7eb..c82659c0 100644
--- a/spring-integration-azure/spring-integration-servicebus/src/main/java/com/microsoft/azure/spring/integration/servicebus/converter/ServiceBusMessageConverter.java
+++ b/spring-integration-azure/spring-integration-servicebus/src/main/java/com/microsoft/azure/spring/integration/servicebus/converter/ServiceBusMessageConverter.java
@@ -6,6 +6,7 @@
 
 package com.microsoft.azure.spring.integration.servicebus.converter;
 
+import com.microsoft.azure.management.Azure;
 import com.microsoft.azure.servicebus.IMessage;
 import com.microsoft.azure.servicebus.Message;
 import com.microsoft.azure.servicebus.MessageBody;
@@ -18,6 +19,8 @@ import org.springframework.util.InvalidMimeTypeException;
 import org.springframework.util.MimeType;
 import org.springframework.util.StringUtils;
 
+import java.time.Duration;
+import java.time.Instant;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
@@ -82,6 +85,11 @@ public class ServiceBusMessageConverter extends AbstractAzureMessageConverter<IM
             serviceBusMessage.setReplyTo(headers.get(MessageHeaders.REPLY_CHANNEL, String.class));
         }
 
+        if (headers.containsKey(AzureHeaders.SCHEDULED_ENQUEUE_MESSAGE)) {
+            serviceBusMessage.setScheduledEnqueueTimeUtc(Instant.now().plus(Duration.ofMillis(
+                    headers.get(AzureHeaders.SCHEDULED_ENQUEUE_MESSAGE, Integer.class))));
+        }
+
         headers.forEach((key, value) -> serviceBusMessage.getProperties().put(key, value.toString()));
     }
 
diff --git a/spring-integration-azure/spring-integration-servicebus/src/test/java/com/microsoft/azure/spring/integration/servicebus/ServiceBusMessageConverterTest.java b/spring-integration-azure/spring-integration-servicebus/src/test/java/com/microsoft/azure/spring/integration/servicebus/ServiceBusMessageConverterTest.java
index b0bc63a1..e7f5903e 100644
--- a/spring-integration-azure/spring-integration-servicebus/src/test/java/com/microsoft/azure/spring/integration/servicebus/ServiceBusMessageConverterTest.java
+++ b/spring-integration-azure/spring-integration-servicebus/src/test/java/com/microsoft/azure/spring/integration/servicebus/ServiceBusMessageConverterTest.java
@@ -17,7 +17,10 @@ import com.microsoft.azure.spring.integration.servicebus.converter.ServiceBusMes
 import com.microsoft.azure.spring.integration.test.support.AzureMessageConverterTest;
 import org.junit.Test;
 import org.springframework.messaging.MessageHeaders;
+import org.springframework.messaging.support.GenericMessage;
+import org.springframework.integration.support.MessageBuilder;
 
+import java.util.HashMap;
 import java.util.List;
 
 import static java.util.Collections.singletonList;
@@ -98,4 +101,14 @@ public class ServiceBusMessageConverterTest extends AzureMessageConverterTest<IM
                 convertedPayload.getPayload(),
                 new ObjectMapper().writeValueAsBytes(internalSequence));
     }
+
+    @Test
+    public void shouldConvertSpringMessageHeaderIntoIMessage() {
+        org.springframework.messaging.Message<String> springMessage =
+                MessageBuilder.withPayload(payload).setHeader("x-delay", 5000).build();
+        IMessage servicebusMessage = getConverter().fromMessage(springMessage, IMessage.class);
+        assertNotNull(servicebusMessage);
+        assertNotNull(servicebusMessage.getScheduledEnqueueTimeUtc());
+
+    }
 }