Merge pull request #28 from olegz/INT-2028

MongoDbMessageStore for M1
2011-08-26 16:28:15 -04:00
parent 96dd1d4d6d 1702413143
commit d895a8d1e1
5 changed files with 551 additions and 37 deletions
--- a/spring-integration-mongodb/src/main/java/org/springframework/integration/mongodb/store/MongoDbMessageStore.java
+++ b/spring-integration-mongodb/src/main/java/org/springframework/integration/mongodb/store/MongoDbMessageStore.java
@@ -19,6 +19,8 @@ package org.springframework.integration.mongodb.store;
 import static org.springframework.data.mongodb.core.query.Criteria.where;

 import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
@@ -35,10 +37,15 @@ import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
 import org.springframework.data.mongodb.core.query.Query;
+import org.springframework.data.mongodb.core.query.Update;
 import org.springframework.integration.Message;
 import org.springframework.integration.MessageHeaders;
 import org.springframework.integration.message.GenericMessage;
+import org.springframework.integration.store.AbstractMessageGroupStore;
+import org.springframework.integration.store.MessageGroup;
+import org.springframework.integration.store.MessageGroupStore;
 import org.springframework.integration.store.MessageStore;
+import org.springframework.integration.store.SimpleMessageGroup;
 import org.springframework.jmx.export.annotation.ManagedAttribute;
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
@@ -47,14 +54,23 @@ import org.springframework.util.StringUtils;
 import com.mongodb.DBObject;

 /**
+ * An implementation of both the {@link MessageStore} and {@link MessageGroupStore}
+ * strategies that relies upon MongoDB for persistence.
+ * 
 * @author Mark Fisher
 * @author Oleg Zhurakousky
 * @since 2.1
 */
-public class MongoDbMessageStore implements MessageStore, BeanClassLoaderAware {
+public class MongoDbMessageStore extends AbstractMessageGroupStore implements MessageStore, BeanClassLoaderAware {

 	private final static String DEFAULT_COLLECTION_NAME = "messages";

+	private final static String GROUP_ID_KEY = "_groupId";
+
+	private final static String MARKED_KEY = "_marked";
+
+	private final static String PAYLOAD_TYPE_KEY = "_payloadType";
+

 	private final MongoTemplate template;

@@ -63,16 +79,22 @@ public class MongoDbMessageStore implements MessageStore, BeanClassLoaderAware {
 	private volatile ClassLoader classLoader = ClassUtils.getDefaultClassLoader();


+	/**
+	 * Create a MongoDbMessageStore using the provided {@link MongoDbFactory}.and the default collection name.
+	 */
 	public MongoDbMessageStore(MongoDbFactory mongoDbFactory) {
 		this(mongoDbFactory, null);
 	}

+	/**
+	 * Create a MongoDbMessageStore using the provided {@link MongoDbFactory} and collection name.
+	 */
 	public MongoDbMessageStore(MongoDbFactory mongoDbFactory, String collectionName) {
 		Assert.notNull(mongoDbFactory, "mongoDbFactory must not be null");
 		MessageReadingMongoConverter converter = new MessageReadingMongoConverter(mongoDbFactory, new MongoMappingContext());
+		converter.afterPropertiesSet();
 		this.template = new MongoTemplate(mongoDbFactory, converter);
 		this.collectionName = (StringUtils.hasText(collectionName)) ? collectionName : DEFAULT_COLLECTION_NAME;
-		//this.template.createCollection(collectionName);
 	}


@@ -82,28 +104,117 @@ public class MongoDbMessageStore implements MessageStore, BeanClassLoaderAware {
 	}

 	public <T> Message<T> addMessage(Message<T> message) {
-		this.template.insert(message, collectionName);
+		Assert.notNull(message, "'message' must not be null");
+		this.template.insert(new MessageWrapper(message, null, false), this.collectionName);
 		return message;
 	}

 	public Message<?> getMessage(UUID id) {
-		return this.template.findOne(this.idQuery(id), Message.class, this.collectionName);
+		Assert.notNull(id, "'id' must not be null");
+		MessageWrapper messageWrapper = this.template.findOne(whereMessageIdIs(id), MessageWrapper.class, this.collectionName);
+		return (messageWrapper != null) ? messageWrapper.getMessage() : null;
 	}

 	@ManagedAttribute
 	public long getMessageCount() {
-		return this.template.getCollection(DEFAULT_COLLECTION_NAME).getCount();
+		return this.template.getCollection(this.collectionName).getCount();
 	}

 	public Message<?> removeMessage(UUID id) {
-		return this.template.findAndRemove(idQuery(id), Message.class, this.collectionName);
+		Assert.notNull(id, "'id' must not be null");
+		MessageWrapper messageWrapper =  this.template.findAndRemove(whereMessageIdIs(id), MessageWrapper.class, this.collectionName);
+		return (messageWrapper != null) ? messageWrapper.getMessage() : null;
 	}

-	private Query idQuery(UUID id) {
-		return new Query(where("_id").is(id.toString()));
+	public MessageGroup getMessageGroup(Object groupId) {
+		Assert.notNull(groupId, "'groupId' must not be null");
+		List<MessageWrapper> messageWrappers = this.template.find(whereGroupIdIs(groupId), MessageWrapper.class, this.collectionName);
+		List<Message<?>> unmarkedMessages = new ArrayList<Message<?>>();
+		List<Message<?>> markedMessages = new ArrayList<Message<?>>();
+		for (MessageWrapper messageWrapper : messageWrappers) {
+			if (messageWrapper.isMarked()) {
+				markedMessages.add(messageWrapper.getMessage());
+			}
+			else {
+				unmarkedMessages.add(messageWrapper.getMessage());
+			}
+		}
+		return new SimpleMessageGroup(unmarkedMessages, markedMessages, groupId, System.currentTimeMillis());
+	}
+
+	public MessageGroup addMessageToGroup(Object groupId, Message<?> message) {
+		Assert.notNull(groupId, "'groupId' must not be null");
+		Assert.notNull(message, "'message' must not be null");
+		MessageWrapper wrapper = new MessageWrapper(message, groupId, false);
+		this.template.insert(wrapper, this.collectionName);
+		return this.getMessageGroup(groupId);
+	}
+
+	public MessageGroup markMessageGroup(MessageGroup group) {
+		Assert.notNull(group, "'group' must not be null");
+		Object groupId = group.getGroupId();
+		List<MessageWrapper> messageWrappers = this.template.find(whereGroupIdIs(groupId), MessageWrapper.class, this.collectionName);
+		for (MessageWrapper messageWrapper : messageWrappers) {
+			this.markMessageFromGroup(groupId, messageWrapper.getMessage());
+		}
+		return this.getMessageGroup(groupId);
+	}
+
+	public MessageGroup removeMessageFromGroup(Object groupId, Message<?> messageToRemove) {
+		Assert.notNull(groupId, "'groupId' must not be null");
+		Assert.notNull(messageToRemove, "'messageToRemove' must not be null");
+		this.removeMessage(messageToRemove.getHeaders().getId());
+		return this.getMessageGroup(groupId);
+	}
+
+	public MessageGroup markMessageFromGroup(Object groupId, Message<?> messageToMark) {
+		Update update = Update.update(MARKED_KEY, true);
+		Query q = whereMessageIdIs(messageToMark.getHeaders().getId());
+		this.template.updateFirst(q, update, this.collectionName);
+		return this.getMessageGroup(groupId);	    
+	}
+
+	public void removeMessageGroup(Object groupId) {
+		List<MessageWrapper> messageWrappers = this.template.find(whereGroupIdIs(groupId), MessageWrapper.class, this.collectionName);
+		for (MessageWrapper messageWrapper : messageWrappers) {
+			this.removeMessageFromGroup(groupId, messageWrapper.getMessage());
+		}
+	}
+
+	@Override
+	public Iterator<MessageGroup> iterator() {
+		List<MessageWrapper> groupedMessages = this.template.find(whereGroupIdExists(), MessageWrapper.class, this.collectionName);
+		Map<Object, MessageGroup> messageGroups = new HashMap<Object, MessageGroup>();
+		for (MessageWrapper groupedMessage : groupedMessages) {
+			Object groupId = groupedMessage.getGroupId();
+			if (!messageGroups.containsKey(groupId)) {
+				messageGroups.put(groupId, this.getMessageGroup(groupId));
+			}
+		}
+		return messageGroups.values().iterator();
 	}


+	/*
+	 * Common Queries
+	 */
+
+	private static Query whereMessageIdIs(UUID id) {
+		return new Query(where("headers.id").is(id.toString()));
+	}
+
+	private static Query whereGroupIdIs(Object groupId) {
+		return new Query(where(GROUP_ID_KEY).is(groupId));
+	}
+
+	private static Query whereGroupIdExists() {
+		return new Query(where(GROUP_ID_KEY).exists(true));
+	}
+
+
+	/**
+	 * Custom implementation of the {@link MappingMongoConverter} strategy.
+	 */
 	private class MessageReadingMongoConverter extends MappingMongoConverter {

 		public MessageReadingMongoConverter(MongoDbFactory mongoDbFactory,
@@ -111,51 +222,68 @@ public class MongoDbMessageStore implements MessageStore, BeanClassLoaderAware {
 			super(mongoDbFactory, mappingContext);
 		}

-
 		@Override
 		public void afterPropertiesSet() {
-			super.afterPropertiesSet();
 			List<Converter<?, ?>> customConverters = new ArrayList<Converter<?,?>>();
 			customConverters.add(new UuidToStringConverter());
 			customConverters.add(new StringToUuidConverter());
 			this.setCustomConversions(new CustomConversions(customConverters));
-		}
-	
-		@Override
-		public void write(Object source, DBObject target) {
-			if (source instanceof Message) {
-				String payloadType = ((Message<?>) source).getPayload().getClass().getName();
-				target.put("_payloadType", payloadType);
-				target.put("_id", ((Message<?>) source).getHeaders().getId().toString());
-			}
-			super.write(source, target);
+			super.afterPropertiesSet();
 		}

+		@Override
+		public void write(Object source, DBObject target) {
+			Message<?> message = null;
+			Object groupId = null;
+			boolean marked = false;
+			if (source instanceof MessageWrapper) {
+				MessageWrapper wrapper = (MessageWrapper) source;
+				message = wrapper.getMessage();
+				groupId = wrapper.getGroupId();
+				marked = wrapper.isMarked();
+			}
+			else {
+				Class<?> sourceType = (source != null) ? source.getClass() : null;
+				throw new IllegalArgumentException("Unexpected source type [" + sourceType + "]. Should be a MessageWrapper.");
+			}
+			target.put(PAYLOAD_TYPE_KEY, message.getPayload().getClass().getName());
+			if (groupId != null) {
+				target.put(GROUP_ID_KEY, groupId);
+			}
+			if (marked) {
+				target.put(MARKED_KEY, marked);
+			}
+			super.write(message, target);
+		}

 		@Override
 		@SuppressWarnings({"unchecked", "rawtypes"})
 		public <S> S read(Class<S> clazz, DBObject source) {
-			if (!Message.class.equals(clazz)) {
+			if (!MessageWrapper.class.equals(clazz)) {
 				return super.read(clazz, source);
 			}
-			Map<String, Object> headers = (Map<String, Object>) source.get("headers");
-			Object payload = source.get("payload");
-			Object payloadType = source.get("_payloadType");
-			if (payloadType != null && payload instanceof DBObject) {
-				try {
-					Class<?> payloadClass = ClassUtils.forName(payloadType.toString(), classLoader);
-					payload = this.read(payloadClass, (DBObject) payload);
-				}
-				catch (Exception e) {
-					throw new IllegalStateException("failed to load class: " + payloadType, e);
+			if (source != null) {
+				Map<String, Object> headers = (Map<String, Object>) source.get("headers");
+				Object payload = source.get("payload");
+				Object payloadType = source.get(PAYLOAD_TYPE_KEY);
+				if (payloadType != null && payload instanceof DBObject) {
+					try {
+						Class<?> payloadClass = ClassUtils.forName(payloadType.toString(), classLoader);
+						payload = this.read(payloadClass, (DBObject) payload);
+					}
+					catch (Exception e) {
+						throw new IllegalStateException("failed to load class: " + payloadType, e);
+					}
 				}
+				GenericMessage message = new GenericMessage(payload, headers);
+				Map innerMap = (Map) new DirectFieldAccessor(message.getHeaders()).getPropertyValue("headers");
+				// using reflection to set ID and TIMESTAMP since they are immutable through MessageHeaders
+				innerMap.put(MessageHeaders.ID, UUID.fromString((String) headers.get(MessageHeaders.ID)));
+				innerMap.put(MessageHeaders.TIMESTAMP, headers.get(MessageHeaders.TIMESTAMP));
+				MessageWrapper wrapper = new MessageWrapper(message, source.get(GROUP_ID_KEY), source.get(MARKED_KEY) != null);
+				return (S) wrapper;
 			}
-			GenericMessage message = new GenericMessage(payload, headers);
-			Map innerMap = (Map) new DirectFieldAccessor(message.getHeaders()).getPropertyValue("headers");
-			// TODO: unpick this mess
-			innerMap.put(MessageHeaders.ID, UUID.fromString(source.get("_id").toString()));
-			innerMap.put(MessageHeaders.TIMESTAMP, headers.get(MessageHeaders.TIMESTAMP));
-			return (S) message;
+			return null;
 		}
 	}

@@ -173,4 +301,35 @@ public class MongoDbMessageStore implements MessageStore, BeanClassLoaderAware {
 		}
 	}

+
+	/**
+	 * Wrapper class used for storing Messages in MongoDB along with their "group" metadata.
+	 */
+	private static final class MessageWrapper {
+
+		private final Object groupId;
+
+		private final boolean marked;
+
+		private final Message<?> message;
+
+		public MessageWrapper(Message<?> message, Object groupId, boolean marked) {
+			this.marked = marked;
+			this.message = message;
+			this.groupId = groupId;
+		}
+
+		public Object getGroupId() {
+			return groupId;
+		}
+
+		public boolean isMarked() {
+			return marked;
+		}
+
+		public Message<?> getMessage() {
+			return message;
+		}
+	}
+
 }
--- a/spring-integration-mongodb/src/test/java/org/springframework/integration/mongodb/rules/MongoDbAvailableTests.java
+++ b/spring-integration-mongodb/src/test/java/org/springframework/integration/mongodb/rules/MongoDbAvailableTests.java
@@ -17,6 +17,11 @@
 package org.springframework.integration.mongodb.rules;

 import org.junit.Rule;
+import org.springframework.data.mongodb.MongoDbFactory;
+import org.springframework.data.mongodb.core.MongoTemplate;
+import org.springframework.data.mongodb.core.SimpleMongoDbFactory;
+
+import com.mongodb.Mongo;

 /**
 * Convenience base class that enables unit test methods to rely upon the {@link MongoDbAvailable} annotation.
@@ -28,5 +33,13 @@ public abstract class MongoDbAvailableTests {

 	@Rule
 	public MongoDbAvailableRule redisAvailableRule = new MongoDbAvailableRule();
+	
+	
+	protected MongoDbFactory prepareMongoFactory() throws Exception{
+		MongoDbFactory mongoDbFactory = new SimpleMongoDbFactory(new Mongo(), "test");
+		MongoTemplate template = new MongoTemplate(mongoDbFactory);
+		template.dropCollection("messages");
+		return mongoDbFactory;
+	}

 }
--- a/spring-integration-mongodb/src/test/java/org/springframework/integration/mongodb/store/MongoDbMessageGroupStoreTests.java
+++ b/spring-integration-mongodb/src/test/java/org/springframework/integration/mongodb/store/MongoDbMessageGroupStoreTests.java
@@ -0,0 +1,302 @@
+/*
+ * Copyright 2007-2011 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.
+ *     You may obtain a copy of the License at
+ *
+ *         http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *     Unless required by applicable law or agreed to in writing, software
+ *     distributed under the License is distributed on an "AS IS" BASIS,
+ *     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *     See the License for the specific language governing permissions and
+ *     limitations under the License.
+ */
+package org.springframework.integration.mongodb.store;
+
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+import java.util.UUID;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+
+import junit.framework.AssertionFailedError;
+
+import org.junit.Test;
+import org.springframework.context.support.ClassPathXmlApplicationContext;
+import org.springframework.data.mongodb.MongoDbFactory;
+import org.springframework.integration.Message;
+import org.springframework.integration.MessageChannel;
+import org.springframework.integration.channel.QueueChannel;
+import org.springframework.integration.message.GenericMessage;
+import org.springframework.integration.mongodb.rules.MongoDbAvailable;
+import org.springframework.integration.mongodb.rules.MongoDbAvailableTests;
+import org.springframework.integration.store.MessageGroup;
+import org.springframework.integration.store.SimpleMessageGroup;
+import org.springframework.integration.support.MessageBuilder;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * @author Oleg Zhurakousky
+ *
+ */
+public class MongoDbMessageGroupStoreTests extends MongoDbAvailableTests {
+
+	@Test
+	@MongoDbAvailable
+	public void testNonExistingEmptyMessageGroup() throws Exception{	
+		MongoDbFactory mongoDbFactory = this.prepareMongoFactory();
+		MongoDbMessageStore store = new MongoDbMessageStore(mongoDbFactory);
+		
+		MessageGroup messageGroup = store.getMessageGroup(1);
+		assertNotNull(messageGroup);
+		assertTrue(messageGroup instanceof SimpleMessageGroup);
+		assertEquals(0, messageGroup.size());
+	}
+	
+	@Test
+	@MongoDbAvailable
+	public void testMessageGroupWithAddedMessage() throws Exception{	
+		MongoDbFactory mongoDbFactory = this.prepareMongoFactory();
+		MongoDbMessageStore store = new MongoDbMessageStore(mongoDbFactory);
+		
+		MessageGroup messageGroup = store.getMessageGroup(1);
+		Message<?> messageA = new GenericMessage<String>("A");
+		Message<?> messageB = new GenericMessage<String>("B");
+		store.addMessageToGroup(1, messageA);
+		messageGroup = store.addMessageToGroup(1, messageB);
+		assertEquals(2, messageGroup.size());
+		Message<?> retrievedMessage = store.getMessage(messageA.getHeaders().getId());
+		assertNotNull(retrievedMessage);
+		assertEquals(retrievedMessage.getHeaders().getId(), messageA.getHeaders().getId());
+		// ensure that 'message_group' header that is only used internally is not propagated 
+		assertNull(retrievedMessage.getHeaders().get("message_group"));
+	}
+	
+	@Test
+	@MongoDbAvailable
+	public void testMessageGroupMarkingMessage() throws Exception{	
+		MongoDbFactory mongoDbFactory = this.prepareMongoFactory();
+		MongoDbMessageStore store = new MongoDbMessageStore(mongoDbFactory);
+		
+		MessageGroup messageGroup = store.getMessageGroup(1);
+		Message<?> messageA = new GenericMessage<String>("A");
+		Message<?> messageB = new GenericMessage<String>("B");
+		store.addMessageToGroup(1, messageA);
+		messageGroup = store.addMessageToGroup(1, messageB);
+		assertEquals(0, messageGroup.getMarked().size());
+		assertEquals(2, messageGroup.getUnmarked().size());
+		
+		messageGroup = store.markMessageFromGroup(1, messageA);
+		assertEquals(1, messageGroup.getMarked().size());
+		assertEquals(1, messageGroup.getUnmarked().size());
+		
+		// validate that the updates were propagated to Mongo as well
+		store = new MongoDbMessageStore(mongoDbFactory);
+		
+		messageGroup = store.getMessageGroup(1);
+		assertEquals(1, messageGroup.getMarked().size());
+		assertEquals(1, messageGroup.getUnmarked().size());
+	}
+	
+	@Test
+	@MongoDbAvailable
+	public void testRemoveMessageGroup() throws Exception{	
+		MongoDbFactory mongoDbFactory = this.prepareMongoFactory();
+		MongoDbMessageStore store = new MongoDbMessageStore(mongoDbFactory);
+		
+		MessageGroup messageGroup = store.getMessageGroup(1);
+		Message<?> message = new GenericMessage<String>("Hello");
+		UUID id = message.getHeaders().getId();
+		messageGroup = store.addMessageToGroup(1, message);
+		assertEquals(1, messageGroup.size());
+		message = store.getMessage(id);
+		assertNotNull(message);
+		
+		store.removeMessageGroup(1);
+		MessageGroup messageGroupA = store.getMessageGroup(1);
+		assertEquals(0, messageGroupA.size());
+		assertFalse(messageGroupA.equals(messageGroup));
+		
+	}
+	
+	@Test
+	@MongoDbAvailable
+	public void testRemoveMessageFromTheGroup() throws Exception{	
+		MongoDbFactory mongoDbFactory = this.prepareMongoFactory();
+		MongoDbMessageStore store = new MongoDbMessageStore(mongoDbFactory);
+		
+		MessageGroup messageGroup = store.getMessageGroup(1);
+		Message<?> message = new GenericMessage<String>("2");
+		store.addMessageToGroup(1, new GenericMessage<String>("1"));
+		store.addMessageToGroup(1, message);
+		messageGroup = store.addMessageToGroup(1, new GenericMessage<String>("3"));
+	
+		assertEquals(3, messageGroup.size());
+		
+		messageGroup = store.removeMessageFromGroup(1, message);
+		assertEquals(2, messageGroup.size());
+	}
+	
+	@Test
+	@MongoDbAvailable
+	public void testMarkAllMessagesInMessageGroup() throws Exception {	
+		MongoDbFactory mongoDbFactory = this.prepareMongoFactory();
+		MongoDbMessageStore store = new MongoDbMessageStore(mongoDbFactory);
+		
+		MessageGroup messageGroup = store.getMessageGroup(1);
+		
+		store.addMessageToGroup(1, new GenericMessage<String>("1"));
+		store.addMessageToGroup(1, new GenericMessage<String>("2"));
+		messageGroup = store.addMessageToGroup(1, new GenericMessage<String>("3"));
+		
+		assertEquals(3, messageGroup.getUnmarked().size());
+		assertEquals(0, messageGroup.getMarked().size());
+		
+		messageGroup = store.markMessageGroup(messageGroup);
+
+		assertEquals(0, messageGroup.getUnmarked().size());
+		assertEquals(3, messageGroup.getMarked().size());
+		
+		store = new MongoDbMessageStore(mongoDbFactory);
+		
+		messageGroup = store.getMessageGroup(1);
+		assertEquals(0, messageGroup.getUnmarked().size());
+		assertEquals(3, messageGroup.getMarked().size());
+	}
+	
+	@Test
+	@MongoDbAvailable
+	public void testMultipleMessageStores() throws Exception{	
+		MongoDbFactory mongoDbFactory = this.prepareMongoFactory();
+		MongoDbMessageStore store1 = new MongoDbMessageStore(mongoDbFactory);
+		MongoDbMessageStore store2 = new MongoDbMessageStore(mongoDbFactory);
+
+		Message<?> message = new GenericMessage<String>("1");
+		store1.addMessageToGroup(1, message);
+		store2.addMessageToGroup(1, new GenericMessage<String>("2"));
+		store1.addMessageToGroup(1, new GenericMessage<String>("3"));
+		
+		MongoDbMessageStore store3 = new MongoDbMessageStore(mongoDbFactory);
+		
+		MessageGroup messageGroup = store3.getMessageGroup(1);
+		
+		assertEquals(3, messageGroup.getUnmarked().size());
+		
+		store3.removeMessageFromGroup(1, message);
+		
+		messageGroup = store2.getMessageGroup(1);
+		assertEquals(2, messageGroup.getUnmarked().size());
+	}
+	
+	@Test
+	@MongoDbAvailable
+	public void testMessageGroupIterator() throws Exception{	
+		MongoDbFactory mongoDbFactory = this.prepareMongoFactory();
+		MongoDbMessageStore store1 = new MongoDbMessageStore(mongoDbFactory);
+		MongoDbMessageStore store2 = new MongoDbMessageStore(mongoDbFactory);
+		
+		Message<?> message = new GenericMessage<String>("1");
+		store2.addMessageToGroup(1, message);
+		store1.addMessageToGroup(2, new GenericMessage<String>("2"));
+		store2.addMessageToGroup(3, new GenericMessage<String>("3"));
+		
+		MongoDbMessageStore store3 = new MongoDbMessageStore(mongoDbFactory);
+		
+		Iterator<MessageGroup> iterator = store3.iterator();
+		int counter = 0;
+		while (iterator.hasNext()) {
+			iterator.next();
+			counter++;
+		}
+		assertEquals(3, counter);
+		
+		store2.removeMessageFromGroup(1, message);
+		
+		iterator = store3.iterator();
+		counter = 0;
+		while (iterator.hasNext()) {
+			iterator.next();
+			counter++;
+		}
+		assertEquals(2, counter);
+	}
+	
+//	@Test
+//	@MongoDbAvailable
+//	public void testConcurrentModifications() throws Exception{	
+//		MongoDbFactory mongoDbFactory = this.prepareMongoFactory();
+//		final MongoDbMessageStore store1 = new MongoDbMessageStore(mongoDbFactory);
+//		final MongoDbMessageStore store2 = new MongoDbMessageStore(mongoDbFactory);
+//		
+//		final Message<?> message = new GenericMessage<String>("1"); 
+//
+//		ExecutorService executor = null;
+//		
+//		final List<Object> failures = new ArrayList<Object>();
+//		
+//		for (int i = 0; i < 100; i++) {
+//			executor = Executors.newCachedThreadPool();
+//			
+//			executor.execute(new Runnable() {	
+//				public void run() {			
+//					MessageGroup group = store1.addMessageToGroup(1, message);
+//					if (group.getUnmarked().size() != 1){
+//						failures.add("ADD");
+//						throw new AssertionFailedError("Failed on ADD");
+//					}	
+//				}
+//			});
+//			executor.execute(new Runnable() {	
+//				public void run() {
+//					MessageGroup group = store2.removeMessageFromGroup(1, message);
+//					if (group.getUnmarked().size() != 0){
+//						failures.add("REMOVE");
+//						throw new AssertionFailedError("Failed on Remove");
+//					}	
+//				}
+//			});
+//			
+//			executor.shutdown();
+//			executor.awaitTermination(10, TimeUnit.SECONDS);
+//			store2.removeMessageFromGroup(1, message); // ensures that if ADD thread executed after REMOVE, the store is empty for the next cycle
+//		}
+//		assertTrue(failures.size() == 0);
+//	}
+		
+	
+	@Test
+	@MongoDbAvailable
+	public void testWithAggregatorWithShutdown() throws Exception{	
+		this.prepareMongoFactory(); // for this test it only ensures that DB was flushed before test
+		
+		ClassPathXmlApplicationContext context = new ClassPathXmlApplicationContext("mongo-aggregator-config.xml", this.getClass());
+		MessageChannel input = context.getBean("inputChannel", MessageChannel.class);
+		QueueChannel output = context.getBean("outputChannel", QueueChannel.class);
+		
+		Message<?> m1 = MessageBuilder.withPayload("1").setSequenceNumber(1).setSequenceSize(3).setCorrelationId(1).build();
+		Message<?> m2 = MessageBuilder.withPayload("2").setSequenceNumber(2).setSequenceSize(3).setCorrelationId(1).build();
+		input.send(m1);
+		assertNull(output.receive(1000));
+		input.send(m2);
+		assertNull(output.receive(1000));
+		context.close();
+		
+		context = new ClassPathXmlApplicationContext("mongo-aggregator-config.xml", this.getClass());
+		input = context.getBean("inputChannel", MessageChannel.class);
+		output = context.getBean("outputChannel", QueueChannel.class);
+		
+		Message<?> m3 = MessageBuilder.withPayload("3").setSequenceNumber(3).setSequenceSize(3).setCorrelationId(1).build();
+		input.send(m3);
+		assertNotNull(output.receive(2000));
+	}
+	
+}
--- a/spring-integration-mongodb/src/test/java/org/springframework/integration/mongodb/store/MongoDbMessageStoreTests.java
+++ b/spring-integration-mongodb/src/test/java/org/springframework/integration/mongodb/store/MongoDbMessageStoreTests.java
@@ -18,6 +18,7 @@ package org.springframework.integration.mongodb.store;

 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;

 import org.junit.Test;

@@ -49,6 +50,19 @@ public class MongoDbMessageStoreTests extends MongoDbAvailableTests{
 		assertEquals(messageToStore.getHeaders(), retrievedMessage.getHeaders());
 		assertEquals(messageToStore, retrievedMessage);
 	}
+	@Test 
+	@MongoDbAvailable
+	public void addThenRemoveWithStringPayload() throws Exception {
+		MongoDbFactory mongoDbFactory = new SimpleMongoDbFactory(new Mongo(), "test");
+		MongoDbMessageStore store = new MongoDbMessageStore(mongoDbFactory);
+		Message<?> messageToStore = MessageBuilder.withPayload("Hello").build();
+		store.addMessage(messageToStore);
+		Message<?> retrievedMessage = store.getMessage(messageToStore.getHeaders().getId());
+		assertNotNull(retrievedMessage);
+		store.removeMessage(retrievedMessage.getHeaders().getId());
+		retrievedMessage = store.getMessage(messageToStore.getHeaders().getId());
+		assertNull(retrievedMessage);
+	}

 	@Test 
 	@MongoDbAvailable
--- a/spring-integration-mongodb/src/test/java/org/springframework/integration/mongodb/store/mongo-aggregator-config.xml
+++ b/spring-integration-mongodb/src/test/java/org/springframework/integration/mongodb/store/mongo-aggregator-config.xml
@@ -0,0 +1,26 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<beans xmlns="http://www.springframework.org/schema/beans"
+	xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+	xmlns:int="http://www.springframework.org/schema/integration"
+	xsi:schemaLocation="http://www.springframework.org/schema/beans http://www.springframework.org/schema/beans/spring-beans.xsd
+		http://www.springframework.org/schema/integration http://www.springframework.org/schema/integration/spring-integration-2.0.xsd">
+
+	
+	<int:aggregator input-channel="inputChannel" output-channel="outputChannel" message-store="mongoStore"/>
+	
+	<int:channel id="outputChannel">
+		<int:queue/>
+	</int:channel>
+	
+	<bean id="mongoStore" class="org.springframework.integration.mongodb.store.MongoDbMessageStore">
+		<constructor-arg ref="mongoConnectionFactory"/>
+	</bean>
+	
+	<bean id="mongoConnectionFactory" class="org.springframework.data.mongodb.core.SimpleMongoDbFactory">
+		<constructor-arg>
+			<bean class="com.mongodb.Mongo"/>
+		</constructor-arg>
+		<constructor-arg value="test"/>
+	</bean>
+
+</beans>