Merge pull request #736 from nickspacek/INT-2920

spring-integration-core/src/main/java/org/springframework/integration/transformer/ClaimCheckOutTransformer.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2011 the original author or authors.
+ * Copyright 2002-2013 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -30,6 +30,7 @@ import org.springframework.util.Assert;
  * 
  * @author Mark Fisher
  * @author Oleg Zhurakousky
+ * @author Nick Spacek
  * @since 2.0
  */
 public class ClaimCheckOutTransformer extends AbstractTransformer {
@@ -56,15 +57,18 @@ public class ClaimCheckOutTransformer extends AbstractTransformer {
 		Assert.notNull(message, "message must not be null");
 		Assert.isTrue(message.getPayload() instanceof UUID, "payload must be a UUID");
 		UUID id = (UUID) message.getPayload();
-		Message<?> retrievedMessage = this.messageStore.getMessage(id);
-		Assert.notNull(retrievedMessage, "unable to locate Message for ID: " + id
-				+ " within MessageStore [" + this.messageStore + "]");
+		Message<?> retrievedMessage;
 		if (this.removeMessage) {
-			this.messageStore.removeMessage(id);
+			retrievedMessage = this.messageStore.removeMessage(id);
 			if (logger.isDebugEnabled()) {
 				logger.debug("Removed Message with claim-check '" + id + "' from the MessageStore.");
 			}
 		}
+		else {
+			retrievedMessage = this.messageStore.getMessage(id);
+		}
+		Assert.notNull(retrievedMessage, "unable to locate Message for ID: " + id
+				+ " within MessageStore [" + this.messageStore + "]");
 		MessageBuilder<?> responseBuilder = MessageBuilder.fromMessage(retrievedMessage);
 		// headers on the 'current' message take precedence
 		responseBuilder.copyHeaders(message.getHeaders());