* Fix MongoDB MS for message removal logic

Author: artembilan

spring-integration-mongodb/src/main/java/org/springframework/integration/mongodb/store/ConfigurableMongoDbMessageStore.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2022 the original author or authors.
+ * Copyright 2013-2023 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.
@@ -101,7 +101,8 @@ public <T> Message<T> addMessage(Message<T> message) {
 	@Override
 	public Message<?> removeMessage(UUID id) {
 		Assert.notNull(id, "'id' must not be null");
-		Query query = Query.query(Criteria.where(MessageDocumentFields.MESSAGE_ID).is(id));
+		Query query = Query.query(Criteria.where(MessageDocumentFields.MESSAGE_ID).is(id)
+				.and(MessageDocumentFields.GROUP_ID).exists(false));
 		MessageDocument document = getMongoTemplate().findAndRemove(query, MessageDocument.class, this.collectionName);
 		return (document != null) ? document.getMessage() : null;
 	}

spring-integration-mongodb/src/main/java/org/springframework/integration/mongodb/store/MongoDbMessageStore.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2023 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.
@@ -258,14 +258,15 @@ public MessageMetadata getMessageMetadata(UUID id) {
 	@Override
 	@ManagedAttribute
 	public long getMessageCount() {
-		return this.template.getCollection(this.collectionName).countDocuments();
+		Query query = Query.query(Criteria.where("headers.id").exists(true).and(GROUP_ID_KEY).exists(false));
+		return this.template.getCollection(this.collectionName).countDocuments(query.getQueryObject());
 	}
 
 	@Override
 	public Message<?> removeMessage(UUID id) {
 		Assert.notNull(id, "'id' must not be null");
-		MessageWrapper messageWrapper =
-				this.template.findAndRemove(whereMessageIdIs(id), MessageWrapper.class, this.collectionName);
+		Query query = Query.query(Criteria.where("headers.id").is(id).and(GROUP_ID_KEY).exists(false));
+		MessageWrapper messageWrapper = this.template.findAndRemove(query, MessageWrapper.class, this.collectionName);
 		return (messageWrapper != null ? messageWrapper.getMessage() : null);
 	}
 

spring-integration-mongodb/src/test/java/org/springframework/integration/mongodb/store/AbstractMongoDbMessageGroupStoreTests.java

@@ -536,6 +536,28 @@ void testWithMessageHistory() {
 				.containsEntry("type", "channel");
 	}
 
+	@Test
+	public void removeMessageDoesntRemoveSameMessageInTheGroup() {
+		GenericMessage<String> testMessage = new GenericMessage<>("test data");
+
+		MessageGroupStore store = getMessageGroupStore();
+
+		store.addMessageToGroup("1", testMessage);
+
+		MessageStore messageStore = (MessageStore) store;
+
+		messageStore.removeMessage(testMessage.getHeaders().getId());
+
+		assertThat(messageStore.getMessageCount()).isEqualTo(0);
+		assertThat(store.getMessageCountForAllMessageGroups()).isEqualTo(1);
+		assertThat(store.messageGroupSize("1")).isEqualTo(1);
+
+		store.removeMessageGroup("1");
+
+		assertThat(store.getMessageCountForAllMessageGroups()).isEqualTo(0);
+		assertThat(store.messageGroupSize("1")).isEqualTo(0);
+	}
+
 	protected abstract MessageGroupStore getMessageGroupStore();
 
 	protected abstract MessageStore getMessageStore();