From b44984f59c893f075e55688a4c3611af447eebb9 Mon Sep 17 00:00:00 2001
From: Oleg Zhurakousky <ozhurakousky@vmware.com>
Date: Wed, 19 Oct 2022 10:51:55 +0200
Subject: [PATCH] GH-929 Ensure AWS Functioininvoker handles Mono<Void> return
 the same way as imperative Consumer

Resolves #929
---
 .../function/adapter/aws/FunctionInvoker.java | 23 +++++++++++--------
 .../adapter/aws/FunctionInvokerTests.java     | 21 +++++++++++++++++
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/FunctionInvoker.java b/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/FunctionInvoker.java
index f8eb05c75..fcdb2bbf2 100644
--- a/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/FunctionInvoker.java
+++ b/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/main/java/org/springframework/cloud/function/adapter/aws/FunctionInvoker.java
@@ -90,28 +90,31 @@ public class FunctionInvoker implements RequestStreamHandler {
 
 	@SuppressWarnings("unchecked")
 	private byte[] buildResult(Message<?> requestMessage, Object output) throws IOException {
-		Message<byte[]> responseMessage;
+		Message<byte[]> responseMessage = null;
 		if (output instanceof Publisher<?>) {
 			List<Object> result = new ArrayList<>();
 			for (Object value : Flux.from((Publisher<?>) output).toIterable()) {
-				if (logger.isInfoEnabled()) {
-					logger.info("Response value: " + value);
+				if (logger.isDebugEnabled()) {
+					logger.debug("Response value: " + value);
 				}
 				result.add(value);
 			}
 			if (result.size() > 1) {
 				output = result;
 			}
-			else {
+			else if (result.size() == 1) {
 				output = result.get(0);
 			}
-
-			if (logger.isInfoEnabled()) {
-				logger.info("OUTPUT: " + output + " - " + output.getClass().getName());
+			else {
+				output = null;
+			}
+			if (output != null) {
+				if (logger.isDebugEnabled()) {
+					logger.debug("OUTPUT: " + output + " - " + output.getClass().getName());
+				}
+				byte[] payload = this.jsonMapper.toJson(output);
+				responseMessage = MessageBuilder.withPayload(payload).build();
 			}
-
-			byte[] payload = this.jsonMapper.toJson(output);
-			responseMessage = MessageBuilder.withPayload(payload).build();
 		}
 		else {
 			responseMessage = (Message<byte[]>) output;
diff --git a/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/test/java/org/springframework/cloud/function/adapter/aws/FunctionInvokerTests.java b/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/test/java/org/springframework/cloud/function/adapter/aws/FunctionInvokerTests.java
index 6d9acc34d..e22bf4b22 100644
--- a/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/test/java/org/springframework/cloud/function/adapter/aws/FunctionInvokerTests.java
+++ b/spring-cloud-function-adapters/spring-cloud-function-adapter-aws/src/test/java/org/springframework/cloud/function/adapter/aws/FunctionInvokerTests.java
@@ -44,6 +44,7 @@ import com.fasterxml.jackson.databind.ObjectMapper;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.cloud.function.json.JsonMapper;
@@ -1024,6 +1025,21 @@ public class FunctionInvokerTests {
 		assertThat(result.get("body")).isEqualTo("\"OK\"");
 	}
 
+	@SuppressWarnings("rawtypes")
+	@Test
+	public void testApiGatewayWithMonoVoidAsReturn() throws Exception {
+		System.setProperty("MAIN_CLASS", ApiGatewayConfiguration.class.getName());
+		System.setProperty("spring.cloud.function.definition", "reactiveWithVoidReturn");
+		FunctionInvoker invoker = new FunctionInvoker();
+
+		InputStream targetStream = new ByteArrayInputStream(this.apiGatewayEvent.getBytes());
+		ByteArrayOutputStream output = new ByteArrayOutputStream();
+		invoker.handleRequest(targetStream, output, null);
+
+		Map result = mapper.readValue(output.toByteArray(), Map.class);
+		assertThat(result.get("body")).isEqualTo("\"OK\"");
+	}
+
 	@Test
 	public void testWithDefaultRoutingFailure() throws Exception {
 		System.setProperty("MAIN_CLASS", SampleConfiguration.class.getName());
@@ -1315,6 +1331,11 @@ public class FunctionInvokerTests {
 			return v -> v.toUpperCase();
 		}
 
+		@Bean
+		public Function<Mono<String>, Mono<Void>> reactiveWithVoidReturn() {
+			return v -> Mono.empty();
+		}
+
 		@Bean
 		public Function<Person, String> uppercasePojo() {
 			return v -> {