Double check that a Function can return a Mono

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

@@ -149,12 +149,25 @@ public class FunctionType {
 				.getType());
 	}
 
-	public FunctionType wrap(Class<?> wrapper) {
-		if (wrapper.isAssignableFrom(getInputWrapper()) || !isWrapper(wrapper)) {
+	public FunctionType wrap(Class<?> input, Class<?> output) {
+		if (!isWrapper(input) && !isWrapper(output)) {
+			return this;
+		}
+		if (!isWrapper(input) || !isWrapper(output)) {
+			throw new IllegalArgumentException("Both wrapper types must be wrappers in ("
+					+ input + ", " + output + ")");
+		}
+		if (input.isAssignableFrom(getInputWrapper())
+				&& output.isAssignableFrom(getOutputWrapper())) {
 			return this;
 		}
 		return new FunctionType(ResolvableType.forClassWithGenerics(Function.class,
-				wrap(wrapper, getInputType()), wrap(wrapper, getOutputType())).getType());
+				wrapper(input, getInputType()), wrapper(output, getOutputType()))
+				.getType());
+	}
+
+	public FunctionType wrap(Class<?> wrapper) {
+		return wrap(wrapper, wrapper);
 	}
 
 	public static FunctionType compose(FunctionType input, FunctionType output) {
@@ -173,7 +186,7 @@ public class FunctionType {
 				.getType());
 	}
 
-	private ResolvableType wrap(Class<?> wrapper, Class<?> type) {
+	private ResolvableType wrapper(Class<?> wrapper, Class<?> type) {
 		return wrap(this, wrapper, type);
 	}
 

spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/config/BeanFactoryFunctionCatalogTests.java

@@ -17,7 +17,9 @@
 package org.springframework.cloud.function.context.config;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -78,6 +80,22 @@ public class BeanFactoryFunctionCatalogTests {
 		assertThat(foos.apply(Flux.just(2)).blockFirst()).isEqualTo("i=2");
 	}
 
+	@Test
+	public void registerFunctionWithMonoType() {
+		processor.register(
+				new FunctionRegistration<Function<Flux<String>, Mono<Map<String, Integer>>>>(
+						flux -> flux.collect(HashMap::new,
+								(map, word) -> map.merge(word, 1, Integer::sum)))
+										.names("foos")
+										.type(FunctionType.from(String.class)
+												.to(Map.class)
+												.wrap(Flux.class, Mono.class).getType()));
+		Function<Flux<String>, Mono<Map<String, Integer>>> foos = processor
+				.lookup(Function.class, "");
+		assertThat(foos.apply(Flux.just("one", "one", "two")).block())
+				.containsEntry("one", 2);
+	}
+
 	@Test
 	public void lookupNonExistentConsumerWithEmptyName() {
 		processor.register(new FunctionRegistration<>(new Foos()).names("foos"));

spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/config/ContextFunctionCatalogAutoConfigurationTests.java

@@ -19,6 +19,7 @@ package org.springframework.cloud.function.context.config;
 import java.net.URL;
 import java.net.URLClassLoader;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Consumer;
@@ -236,6 +237,22 @@ public class ContextFunctionCatalogAutoConfigurationTests {
 				.isAssignableFrom(Publisher.class);
 	}
 
+	@Test
+	public void monoFunction() {
+		create(MonoConfiguration.class);
+		assertThat(context.getBean("function")).isInstanceOf(Function.class);
+		assertThat(catalog.<Function<?, ?>>lookup(Function.class, "function"))
+				.isInstanceOf(Function.class);
+		assertThat(inspector.isMessage(catalog.lookup(Function.class, "function")))
+				.isFalse();
+		assertThat(inspector.getInputType(catalog.lookup(Function.class, "function")))
+				.isAssignableFrom(String.class);
+		assertThat(inspector.getInputWrapper(catalog.lookup(Function.class, "function")))
+				.isAssignableFrom(Flux.class);
+		assertThat(inspector.getOutputWrapper(catalog.lookup(Function.class, "function")))
+				.isAssignableFrom(Mono.class);
+	}
+
 	@Test
 	public void messageFunction() {
 		create(MessageConfiguration.class);
@@ -756,6 +773,16 @@ public class ContextFunctionCatalogAutoConfigurationTests {
 		}
 	}
 
+	@EnableAutoConfiguration
+	@Configuration
+	protected static class MonoConfiguration {
+		@Bean
+		public Function<Flux<String>, Mono<Map<String, Integer>>> function() {
+			return flux -> flux.collect(HashMap::new,
+					(map, word) -> map.merge(word, 1, Integer::sum));
+		}
+	}
+
 	@EnableAutoConfiguration
 	@Configuration
 	protected static class MessageConfiguration {