diff --git a/spring-batch-integration/src/main/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandler.java b/spring-batch-integration/src/main/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandler.java
index e84e2507b..e57cbb076 100644
--- a/spring-batch-integration/src/main/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandler.java
+++ b/spring-batch-integration/src/main/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandler.java
@@ -11,6 +11,7 @@ import org.springframework.batch.core.StepExecution;
import org.springframework.batch.core.partition.PartitionHandler;
import org.springframework.batch.core.partition.StepExecutionSplitter;
import org.springframework.batch.core.repository.JobRepository;
+import org.springframework.integration.MessageTimeoutException;
import org.springframework.integration.annotation.Aggregator;
import org.springframework.integration.annotation.MessageEndpoint;
import org.springframework.integration.annotation.Payloads;
@@ -33,6 +34,7 @@ import org.springframework.util.Assert;
*
* @author Dave Syer
* @author Will Schipp
+ * @author Michael Minella
*
*/
@MessageEndpoint
@@ -106,10 +108,10 @@ public class MessageChannelPartitionHandler implements PartitionHandler {
}
/**
- * Sends {@link StepExecutionRequest} objects to the request channel of the {@link MessagingOperations}, and then
+ * Sends {@link StepExecutionRequest} objects to the request channel of the {@link MessagingTemplate}, and then
* receives the result back as a list of {@link StepExecution} on a reply channel. Use the {@link #aggregate(List)}
* method as an aggregator of the individual remote replies. The receive timeout needs to be set realistically in
- * the {@link MessagingOperations} and the aggregator, so that there is a good chance of all work being done.
+ * the {@link MessagingTemplate} and the aggregator, so that there is a good chance of all work being done.
*
* @see PartitionHandler#handle(StepExecutionSplitter, StepExecution)
*/
@@ -134,9 +136,13 @@ public class MessageChannelPartitionHandler implements PartitionHandler {
@SuppressWarnings("unchecked")
Message> message = (Message>) messagingGateway.receive(replyChannel);
- if (logger.isDebugEnabled()) {
+
+ if(message == null) {
+ throw new MessageTimeoutException("Timeout occurred before all partitions returned");
+ } else if (logger.isDebugEnabled()) {
logger.debug("Received replies: " + message);
}
+
Collection result = message.getPayload();
return result;
diff --git a/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTest.java b/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTest.java
index 3b5d812c4..d382a1df2 100644
--- a/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTest.java
+++ b/spring-batch-integration/src/test/java/org/springframework/batch/integration/partition/MessageChannelPartitionHandlerTest.java
@@ -1,24 +1,26 @@
package org.springframework.batch.integration.partition;
+import org.junit.Test;
+import org.springframework.batch.core.StepExecution;
+import org.springframework.batch.core.partition.StepExecutionSplitter;
+import org.springframework.integration.MessageTimeoutException;
+import org.springframework.integration.core.MessagingTemplate;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.PollableChannel;
+
+import java.util.Collection;
+import java.util.Collections;
+
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.anyObject;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
-import java.util.Collection;
-import java.util.Collections;
-
-import org.junit.Test;
-import org.springframework.batch.core.StepExecution;
-import org.springframework.batch.core.partition.StepExecutionSplitter;
-import org.springframework.integration.core.MessagingTemplate;
-import org.springframework.messaging.Message;
-import org.springframework.messaging.PollableChannel;
-
/**
*
* @author Will Schipp
+ * @author Michael Minella
*
*/
public class MessageChannelPartitionHandlerTest {
@@ -72,4 +74,21 @@ public class MessageChannelPartitionHandlerTest {
}
+ @Test(expected = MessageTimeoutException.class)
+ public void messageReceiveTimeout() throws Exception {
+ //execute with no default set
+ messageChannelPartitionHandler = new MessageChannelPartitionHandler();
+ //mock
+ StepExecution masterStepExecution = mock(StepExecution.class);
+ StepExecutionSplitter stepExecutionSplitter = mock(StepExecutionSplitter.class);
+ MessagingTemplate operations = mock(MessagingTemplate.class);
+ Message message = mock(Message.class);
+ //when
+ when(message.getPayload()).thenReturn(Collections.emptyList());
+ //set
+ messageChannelPartitionHandler.setMessagingOperations(operations);
+
+ //execute
+ Collection executions = messageChannelPartitionHandler.handle(stepExecutionSplitter, masterStepExecution);
+ }
}