* Add FluxAggregatorMessageHandlerTests

* Fix `FluxAggregatorMessageHandler` for the default `messageForWindowFlux`
to rely on the `publish().autoConnect()` to build a target `Flux` for
window and also copy headers from the first `Message` in the window.
Looks like `switchOnFirst()` doesn't work somehow with windows or I just
don't use it properly, although it isn't clear how to continue chain but
get the whole `Flux` as a single entry for message payload

Author: artembilan

spring-integration-core/src/main/java/org/springframework/integration/aggregator/FluxAggregatorMessageHandler.java

@@ -165,20 +165,14 @@ protected void handleMessageInternal(Message<?> message) {
 	}
 
 	private Mono<Message<?>> messageForWindowFlux(Flux<Message<?>> messageFlux) {
-		return messageFlux
-				.switchOnFirst((signal, window) -> {
-					if (signal.hasValue()) {
-						return Flux.<Message<?>>just(
-								getMessageBuilderFactory()
-										.withPayload(window)
-										.copyHeaders(signal.get().getHeaders())
-										.build());
-					}
-					else {
-						return window.thenMany(Flux.empty());
-					}
-				})
-				.next();
+		Flux<Message<?>> window = messageFlux.publish().autoConnect();
+		return window
+				.next()
+				.map((first) ->
+						getMessageBuilderFactory()
+								.withPayload(Flux.concat(Mono.just(first), window))
+								.copyHeaders(first.getHeaders())
+								.build());
 	}
 
 	private static Integer sequenceSizeHeader(Message<?> message) {

spring-integration-core/src/test/java/org/springframework/integration/aggregator/FluxAggregatorMessageHandlerTests.java

@@ -18,23 +18,31 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
+import java.time.Duration;
+import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.Executors;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
 import org.junit.jupiter.api.Test;
 
 import org.springframework.integration.IntegrationMessageHeaderAccessor;
 import org.springframework.integration.channel.QueueChannel;
 import org.springframework.integration.support.MessageBuilder;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageHeaders;
+import org.springframework.messaging.support.GenericMessage;
 
 import reactor.core.publisher.Flux;
 import reactor.test.StepVerifier;
-import reactor.util.Loggers;
 
 /**
  * @author Artem Bilan
  *
  * @since 5.2
  */
+@SuppressWarnings("unchecked")
 class FluxAggregatorMessageHandlerTests {
 
 	@Test
@@ -62,23 +70,176 @@ void testDefaultAggregation() {
 		Object payload = result.getPayload();
 		assertThat(payload).isInstanceOf(Flux.class);
 
-		Loggers.useVerboseConsoleLoggers();
-
-		@SuppressWarnings("unchecked")
 		Flux<Message<?>> window = (Flux<Message<?>>) payload;
 
 		StepVerifier.create(
 				window.map(Message::getPayload)
-						.cast(String.class)
-						.log())
-				.expectNext("0", "2", "4", "6", "8", "10", "12", "14", "16", "18", "20")
-				/*.expectNextSequence(
+						.cast(String.class))
+				.expectNextSequence(
 						IntStream.iterate(0, i -> i + 2)
 								.limit(10)
 								.mapToObj(Objects::toString)
-								.collect(Collectors.toList()))*/
-				.thenCancel()
-				.verify();
+								.collect(Collectors.toList()))
+				.verifyComplete();
+
+		result = resultChannel.receive(10_000);
+		assertThat(result).isNotNull()
+				.extracting(Message::getHeaders)
+				.satisfies((headers) ->
+						assertThat((MessageHeaders) headers)
+								.containsEntry(IntegrationMessageHeaderAccessor.CORRELATION_ID, 1));
+
+		payload = result.getPayload();
+		window = (Flux<Message<?>>) payload;
+
+		StepVerifier.create(
+				window.map(Message::getPayload)
+						.cast(String.class))
+				.expectNextSequence(
+						IntStream.iterate(1, i -> i + 2)
+								.limit(10)
+								.mapToObj(Objects::toString)
+								.collect(Collectors.toList()))
+				.verifyComplete();
+	}
+
+	@Test
+	void testCustomCombineFunction() {
+		QueueChannel resultChannel = new QueueChannel();
+		FluxAggregatorMessageHandler fluxAggregatorMessageHandler = new FluxAggregatorMessageHandler();
+		fluxAggregatorMessageHandler.setOutputChannel(resultChannel);
+		fluxAggregatorMessageHandler.setWindowSize(10);
+		fluxAggregatorMessageHandler.setCombineFunction(
+				(messageFlux) ->
+						messageFlux
+								.map(Message::getPayload)
+								.collectList()
+								.map(GenericMessage::new));
+
+		for (int i = 0; i < 20; i++) {
+			Message<?> messageToAggregate =
+					MessageBuilder.withPayload(i)
+							.setCorrelationId(i % 2)
+							.build();
+			fluxAggregatorMessageHandler.handleMessage(messageToAggregate);
+		}
+
+		Message<?> result = resultChannel.receive(10_000);
+		assertThat(result).isNotNull();
+
+		Object payload = result.getPayload();
+		assertThat(payload)
+				.isInstanceOf(List.class)
+				.asList()
+				.containsExactly(
+						IntStream.iterate(0, i -> i + 2)
+								.limit(10)
+								.boxed()
+								.toArray());
+
+		result = resultChannel.receive(10_000);
+		assertThat(result).isNotNull();
+
+		payload = result.getPayload();
+		assertThat(payload)
+				.isInstanceOf(List.class)
+				.asList()
+				.containsExactly(
+						IntStream.iterate(1, i -> i + 2)
+								.limit(10)
+								.boxed()
+								.toArray());
+	}
+
+	@Test
+	void testWindowTimespan() {
+		QueueChannel resultChannel = new QueueChannel();
+		FluxAggregatorMessageHandler fluxAggregatorMessageHandler = new FluxAggregatorMessageHandler();
+		fluxAggregatorMessageHandler.setOutputChannel(resultChannel);
+		fluxAggregatorMessageHandler.setWindowTimespan(Duration.ofMillis(100));
+
+		Executors.newSingleThreadExecutor()
+				.submit(() -> {
+					for (int i = 0; i < 10; i++) {
+						Message<?> messageToAggregate =
+								MessageBuilder.withPayload(i)
+										.setCorrelationId("1")
+										.build();
+						fluxAggregatorMessageHandler.handleMessage(messageToAggregate);
+						Thread.sleep(20);
+					}
+					return null;
+				});
+
+		Message<?> result = resultChannel.receive(10_000);
+		assertThat(result).isNotNull();
+
+		Flux<Message<?>> window = (Flux<Message<?>>) result.getPayload();
+
+		List<Integer> messageList =
+				window.map(Message::getPayload)
+						.cast(Integer.class)
+						.collectList()
+						.block(Duration.ofSeconds(10));
+
+		assertThat(messageList)
+				.isNotEmpty()
+				.hasSizeLessThan(10)
+				.contains(0, 1);
+
+		result = resultChannel.receive(10_000);
+		assertThat(result).isNotNull();
+
+		window = (Flux<Message<?>>) result.getPayload();
+
+		messageList =
+				window.map(Message::getPayload)
+						.cast(Integer.class)
+						.collectList()
+						.block(Duration.ofSeconds(10));
+
+		assertThat(messageList)
+				.isNotEmpty()
+				.hasSizeLessThan(10)
+				.doesNotContain(0, 1);
+	}
+
+	@Test
+	void testBoundaryTrigger() {
+		QueueChannel resultChannel = new QueueChannel();
+		FluxAggregatorMessageHandler fluxAggregatorMessageHandler = new FluxAggregatorMessageHandler();
+		fluxAggregatorMessageHandler.setOutputChannel(resultChannel);
+		fluxAggregatorMessageHandler.setBoundaryTrigger((message -> "terminate".equals(message.getPayload())));
+
+		for (int i = 0; i < 3; i++) {
+			Message<?> messageToAggregate =
+					MessageBuilder.withPayload("" + i)
+							.setCorrelationId("1")
+							.build();
+			fluxAggregatorMessageHandler.handleMessage(messageToAggregate);
+		}
+
+		fluxAggregatorMessageHandler.handleMessage(
+				MessageBuilder.withPayload("terminate")
+						.setCorrelationId("1")
+						.build());
+
+		fluxAggregatorMessageHandler.handleMessage(
+				MessageBuilder.withPayload("next")
+						.setCorrelationId("1")
+						.build());
+
+		Message<?> result = resultChannel.receive(10_000);
+		assertThat(result).isNotNull();
+
+		Flux<Message<?>> window = (Flux<Message<?>>) result.getPayload();
+
+		StepVerifier.create(
+				window.map(Message::getPayload)
+						.cast(String.class))
+				.expectNext("0", "1", "2")
+				.expectNext("terminate")
+				.verifyComplete();
 	}
 
 }