Add support for composing function of Message with plain function

Fixes gh-267 at least for the most common cases.

spring-cloud-function-context/pom.xml

@@ -53,6 +53,11 @@
 			<artifactId>spring-cloud-function-compiler</artifactId>
 			<scope>test</scope>
 		</dependency>
+		<dependency>
+			<groupId>io.projectreactor</groupId>
+			<artifactId>reactor-test</artifactId>
+			<scope>test</scope>
+		</dependency>
 	</dependencies>
 	<build>
 		<plugins>

spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/AbstractComposableFunctionRegistry.java

@@ -362,6 +362,12 @@ public abstract class AbstractComposableFunctionRegistry implements FunctionRegi
 		FunctionType bType = bReg.getType();
 		Object a = aReg.getTarget();
 		Object b = bReg.getTarget();
+		if (aType != null && bType != null) {
+			if (aType.isMessage() && !bType.isMessage()) {
+				bType = bType.message();
+				b = message(b);
+			}
+		}
 		Object composedFunction = null;
 		if (a instanceof Supplier && b instanceof Function) {
 			Supplier<Flux<Object>> supplier = (Supplier<Flux<Object>>) a;
@@ -420,6 +426,19 @@ public abstract class AbstractComposableFunctionRegistry implements FunctionRegi
 				.type(FunctionType.compose(aType, bType));
 	}
 
+	private Object message(Object input) {
+		if (input instanceof Supplier) {
+			return new MessageSupplier((Supplier<?>) input);
+		}
+		if (input instanceof Consumer) {
+			return new MessageConsumer((Consumer<?>) input);
+		}
+		if (input instanceof Function) {
+			return new MessageFunction((Function<?, ?>) input);
+		}
+		return input;
+	}
+
 	@SuppressWarnings("unchecked")
 	private <T> T doLookup(Class<?> type, String name) {
 		T function = null;

spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/MessageConsumer.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright 2019-2019 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.cloud.function.context.catalog;
+
+import java.util.function.Consumer;
+
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+
+import org.springframework.messaging.Message;
+
+/**
+ * @author Dave Syer
+ */
+public class MessageConsumer implements Consumer<Publisher<Message<?>>> {
+
+	private Consumer<Object> delegate;
+
+	@SuppressWarnings("unchecked")
+	public MessageConsumer(Consumer<?> input) {
+		this.delegate = (Consumer<Object>) input;
+	}
+
+	@Override
+	public void accept(Publisher<Message<?>> input) {
+		Flux.from(input).map(Message::getPayload).subscribe(this.delegate::accept);
+	}
+
+}

spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/MessageFunction.java

@@ -0,0 +1,95 @@
+/*
+ * Copyright 2019-2019 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.cloud.function.context.catalog;
+
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+import org.springframework.cloud.function.core.FluxConsumer;
+import org.springframework.cloud.function.core.FluxFunction;
+import org.springframework.cloud.function.core.FluxToMonoFunction;
+import org.springframework.cloud.function.core.MonoToFluxFunction;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageHeaders;
+import org.springframework.messaging.support.MessageBuilder;
+
+/**
+ * @author Dave Syer
+ */
+public class MessageFunction
+		implements Function<Publisher<Message<?>>, Publisher<Message<?>>> {
+
+	private Function<?, ?> delegate;
+
+	public MessageFunction(Function<?, ?> delegate) {
+		this.delegate = delegate;
+	}
+
+	@Override
+	public Publisher<Message<?>> apply(Publisher<Message<?>> input) {
+		Flux<Message<?>> flux = Flux.from(input);
+		if (this.delegate instanceof FluxFunction) {
+			@SuppressWarnings("unchecked")
+			Function<Object, Object> target = (Function<Object, Object>) ((FluxFunction<?, ?>) this.delegate)
+					.getTarget();
+			return flux.map(
+					value -> MessageBuilder.withPayload(target.apply(value.getPayload()))
+							.copyHeaders(value.getHeaders()).build());
+		}
+		if (this.delegate instanceof MonoToFluxFunction) {
+			@SuppressWarnings("unchecked")
+			Function<Mono<Object>, Flux<Object>> target = ((MonoToFluxFunction<Object, Object>) this.delegate)
+					.getTarget();
+			return flux.next()
+					.flatMapMany(value -> target.apply(Mono.just(value.getPayload()))
+							.map(object -> MessageBuilder.withPayload(object)
+									.copyHeaders(value.getHeaders()).build()));
+		}
+		if (this.delegate instanceof FluxToMonoFunction) {
+			@SuppressWarnings("unchecked")
+			Function<Flux<Object>, Mono<Object>> target = ((FluxToMonoFunction<Object, Object>) this.delegate)
+					.getTarget();
+			AtomicReference<MessageHeaders> headers = new AtomicReference<>();
+			return target.apply(flux.map(messsage -> {
+				headers.set(messsage.getHeaders());
+				return messsage.getPayload();
+			})).map(payload -> MessageBuilder.withPayload(payload)
+					.copyHeaders(headers.get()).build());
+		}
+		if (this.delegate instanceof FluxConsumer) {
+			@SuppressWarnings("unchecked")
+			FluxConsumer<Object> target = ((FluxConsumer<Object>) this.delegate);
+			AtomicReference<MessageHeaders> headers = new AtomicReference<>();
+			Mono<Void> mapped = target.apply(flux.map(messsage -> {
+				headers.set(messsage.getHeaders());
+				return messsage.getPayload();
+			}));
+			return mapped.map(value -> MessageBuilder.createMessage(null, headers.get()));
+		}
+		// TODO: cover the case that delegate is actually Function<Flux,Flux>
+		@SuppressWarnings("unchecked")
+		Function<Object, Object> function = (Function<Object, Object>) this.delegate;
+		return flux.map(
+				value -> MessageBuilder.withPayload(function.apply(value.getPayload()))
+						.copyHeaders(value.getHeaders()).build());
+	}
+
+}

spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/MessageSupplier.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2019-2019 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.cloud.function.context.catalog;
+
+import java.util.function.Supplier;
+
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+import org.springframework.cloud.function.core.FluxSupplier;
+import org.springframework.cloud.function.core.MonoSupplier;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.support.MessageBuilder;
+
+/**
+ * @author Dave Syer
+ */
+public class MessageSupplier implements Supplier<Publisher<Message<?>>> {
+
+	private Supplier<?> delegate;
+
+	public MessageSupplier(Supplier<?> delegate) {
+		this.delegate = delegate;
+	}
+
+	@Override
+	public Publisher<Message<?>> get() {
+		if (this.delegate instanceof FluxSupplier) {
+			return ((Flux<?>) this.delegate.get())
+					.map(value -> MessageBuilder.withPayload(value).build());
+		}
+		if (this.delegate instanceof MonoSupplier) {
+			return ((Mono<?>) this.delegate.get())
+					.map(value -> MessageBuilder.withPayload(value).build());
+		}
+		Object product = this.delegate.get();
+		if (product instanceof Publisher) {
+			return Flux.from((Publisher<?>) product)
+					.map(value -> MessageBuilder.withPayload(value).build());
+		}
+		if (product instanceof Iterable) {
+			return Flux.fromIterable((Iterable<?>) product)
+					.map(value -> MessageBuilder.withPayload(value).build());
+		}
+		return Mono.just(MessageBuilder.withPayload(product).build());
+	}
+
+}

spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/InMemoryFunctionCatalogTests.java

@@ -18,7 +18,6 @@ package org.springframework.cloud.function.context.catalog;
 
 import java.util.function.Function;
 
-import org.junit.Ignore;
 import org.junit.Test;
 import reactor.core.publisher.Flux;
 
@@ -111,12 +110,10 @@ public class InMemoryFunctionCatalogTests {
 	}
 
 	@Test
-	@Ignore
 	public void testFunctionCompositionMixedMessages() {
 		FunctionRegistration<UpperCaseMessage> upperCaseRegistration = new FunctionRegistration<>(
 				new UpperCaseMessage(), "uppercase")
 						.type(FunctionType.of(UpperCaseMessage.class).getType());
-		// TODO: make this work with plain Reverse (not message)
 		FunctionRegistration<Reverse> reverseRegistration = new FunctionRegistration<>(
 				new Reverse(), "reverse").type(FunctionType.of(Reverse.class).getType());
 		InMemoryFunctionCatalog catalog = new InMemoryFunctionCatalog();
@@ -128,9 +125,11 @@ public class InMemoryFunctionCatalogTests {
 		assertThat(catalog.getFunctionType("uppercase|reverse").isMessage()).isTrue();
 
 		assertThat(lookedUpFunction).isNotNull();
-		assertThat(lookedUpFunction
-				.apply(Flux.just(MessageBuilder.withPayload("star").build())).blockFirst()
-				.getPayload()).isEqualTo("RATS");
+		Message<String> message = lookedUpFunction.apply(Flux
+				.just(MessageBuilder.withPayload("star").setHeader("foo", "bar").build()))
+				.blockFirst();
+		assertThat(message.getPayload()).isEqualTo("RATS");
+		assertThat(message.getHeaders().get("foo")).isEqualTo("bar");
 	}
 
 	private static class UpperCase implements Function<String, String> {

spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/MessageConsumerTests.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2019-2019 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.cloud.function.context.catalog;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Consumer;
+
+import org.junit.Test;
+import reactor.core.publisher.Flux;
+
+import org.springframework.messaging.support.MessageBuilder;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * @author Dave Syer
+ */
+public class MessageConsumerTests {
+
+	private List<String> items = new ArrayList<>();
+
+	@Test
+	public void plainConsumer() {
+		MessageConsumer consumer = new MessageConsumer(input());
+		consumer.accept(Flux
+				.just(MessageBuilder.withPayload("foo").setHeader("foo", "bar").build()));
+		assertThat(this.items).hasSize(1);
+	}
+
+	private Consumer<String> input() {
+		return value -> this.items.add(value);
+	}
+
+}

spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/MessageFunctionTests.java

@@ -0,0 +1,111 @@
+/*
+ * Copyright 2019-2019 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.cloud.function.context.catalog;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+import org.junit.Test;
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+import reactor.test.StepVerifier;
+
+import org.springframework.cloud.function.core.FluxConsumer;
+import org.springframework.cloud.function.core.FluxFunction;
+import org.springframework.cloud.function.core.FluxToMonoFunction;
+import org.springframework.cloud.function.core.MonoToFluxFunction;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.support.MessageBuilder;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * @author Dave Syer
+ */
+public class MessageFunctionTests {
+
+	private List<String> items = new ArrayList<>();
+
+	@Test
+	public void plainFunction() {
+		MessageFunction function = new MessageFunction(uppercase());
+		Publisher<Message<?>> result = function.apply(Flux
+				.just(MessageBuilder.withPayload("foo").setHeader("foo", "bar").build()));
+		StepVerifier.create(result).assertNext(message -> {
+			assertThat(message.getPayload()).isEqualTo("FOO");
+			assertThat(message.getHeaders()).containsEntry("foo", "bar");
+		});
+	}
+
+	@Test
+	public void fluxFunction() {
+		MessageFunction function = new MessageFunction(new FluxFunction<>(uppercase()));
+		Publisher<Message<?>> result = function.apply(Flux
+				.just(MessageBuilder.withPayload("foo").setHeader("foo", "bar").build()));
+		StepVerifier.create(result).assertNext(message -> {
+			assertThat(message.getPayload()).isEqualTo("FOO");
+			assertThat(message.getHeaders()).containsEntry("foo", "bar");
+		});
+	}
+
+	@Test
+	public void fluxToMonoFunction() {
+		MessageFunction function = new MessageFunction(
+				new FluxToMonoFunction<>(flux -> flux.next().map(uppercase())));
+		Publisher<Message<?>> result = function.apply(Flux
+				.just(MessageBuilder.withPayload("foo").setHeader("foo", "bar").build()));
+		StepVerifier.create(result).assertNext(message -> {
+			assertThat(message.getPayload()).isEqualTo("FOO");
+			assertThat(message.getHeaders()).containsEntry("foo", "bar");
+		});
+	}
+
+	@Test
+	public void monoToFunction() {
+		MessageFunction function = new MessageFunction(
+				new MonoToFluxFunction<>(mono -> Flux.from(mono.map(uppercase()))));
+		Publisher<Message<?>> result = function.apply(Flux
+				.just(MessageBuilder.withPayload("foo").setHeader("foo", "bar").build()));
+		StepVerifier.create(result).assertNext(message -> {
+			assertThat(message.getPayload()).isEqualTo("FOO");
+			assertThat(message.getHeaders()).containsEntry("foo", "bar");
+		});
+	}
+
+	@Test
+	public void fluxConsumer() {
+		MessageFunction function = new MessageFunction(new FluxConsumer<>(stash()));
+		Publisher<Message<?>> result = function.apply(Flux
+				.just(MessageBuilder.withPayload("foo").setHeader("foo", "bar").build()));
+		StepVerifier.create(result).assertNext(message -> {
+			assertThat(message.getPayload()).isEqualTo(null);
+			assertThat(message.getHeaders()).containsEntry("foo", "bar");
+			assertThat(this.items).hasSize(1);
+		});
+	}
+
+	private Consumer<String> stash() {
+		return value -> this.items.add(value);
+	}
+
+	private Function<String, String> uppercase() {
+		return value -> value.toUpperCase();
+	}
+
+}

spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/MessageSupplierTests.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2019-2019 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.cloud.function.context.catalog;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.function.Supplier;
+
+import org.junit.Test;
+import reactor.core.publisher.Flux;
+import reactor.test.StepVerifier;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * @author Dave Syer
+ */
+public class MessageSupplierTests {
+
+	@Test
+	public void plainSupplier() {
+		MessageSupplier supplier = new MessageSupplier(input());
+		StepVerifier.create(supplier.get()).assertNext(message -> {
+			assertThat(message.getPayload()).isEqualTo("foo");
+			assertThat(message.getHeaders()).isEmpty();
+		});
+	}
+
+	@Test
+	public void collectionSupplier() {
+		MessageSupplier supplier = new MessageSupplier(inputs());
+		StepVerifier.create(supplier.get()).assertNext(message -> {
+			assertThat(message.getPayload()).isEqualTo("foo");
+			assertThat(message.getHeaders()).isEmpty();
+		}).assertNext(message -> {
+			assertThat(message.getPayload()).isEqualTo("bar");
+			assertThat(message.getHeaders()).isEmpty();
+		});
+	}
+
+	@Test
+	public void fluxSupplier() {
+		MessageSupplier supplier = new MessageSupplier(flux());
+		StepVerifier.create(supplier.get()).assertNext(message -> {
+			assertThat(message.getPayload()).isEqualTo("foo");
+			assertThat(message.getHeaders()).isEmpty();
+		}).assertNext(message -> {
+			assertThat(message.getPayload()).isEqualTo("bar");
+			assertThat(message.getHeaders()).isEmpty();
+		});
+	}
+
+	private Supplier<String> input() {
+		return () -> "foo";
+	}
+
+	private Supplier<Collection<String>> inputs() {
+		return () -> Arrays.asList("foo", "bar");
+	}
+
+	private Supplier<Flux<String>> flux() {
+		return () -> Flux.just("foo", "bar");
+	}
+
+}