From ed6642e18e0e12962c556bad6803ebf4e3dfdedc Mon Sep 17 00:00:00 2001
From: Oleg Zhurakousky <ozhurakousky@vmware.com>
Date: Mon, 21 Feb 2022 14:50:29 +0100
Subject: [PATCH] GH-816 Fix support for function that returns Iterable with
 target-protocol set

Resolves #816
---
 .../function/context/catalog/FunctionTypeUtils.java |  2 +-
 .../context/catalog/SimpleFunctionRegistry.java     |  9 ++++++++-
 .../BeanFactoryAwareFunctionRegistryTests.java      | 13 +++++++++++++
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/FunctionTypeUtils.java b/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/FunctionTypeUtils.java
index 0574ae627..ed5c2959f 100644
--- a/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/FunctionTypeUtils.java
+++ b/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/FunctionTypeUtils.java
@@ -96,7 +96,7 @@ public final class FunctionTypeUtils {
 		}
 		type = getGenericType(type);
 		Class<?> rawType = type instanceof ParameterizedType ? getRawType(type) : (Class<?>) type;
-		return Collection.class.isAssignableFrom(rawType);
+		return Iterable.class.isAssignableFrom(rawType);
 	}
 
 	public static boolean isTypeArray(Type type) {
diff --git a/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/SimpleFunctionRegistry.java b/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/SimpleFunctionRegistry.java
index 1bfa88fd5..d942a885c 100644
--- a/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/SimpleFunctionRegistry.java
+++ b/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/SimpleFunctionRegistry.java
@@ -53,6 +53,7 @@ import org.springframework.cloud.function.context.FunctionProperties.FunctionCon
 import org.springframework.cloud.function.context.FunctionRegistration;
 import org.springframework.cloud.function.context.FunctionRegistry;
 import org.springframework.cloud.function.context.config.RoutingFunction;
+import org.springframework.cloud.function.context.message.MessageUtils;
 import org.springframework.cloud.function.core.FunctionInvocationHelper;
 import org.springframework.cloud.function.json.JsonMapper;
 import org.springframework.context.expression.BeanFactoryResolver;
@@ -1142,7 +1143,13 @@ public class SimpleFunctionRegistry implements FunctionRegistry {
 				convertedOutput = this.convertOutputPublisherIfNecessary((Publisher) output, type, contentType);
 			}
 			else if (output instanceof Message) {
-				convertedOutput = this.convertOutputMessageIfNecessary(output, ObjectUtils.isEmpty(contentType) ? null : contentType[0]);
+				Message m = (Message) output;
+				if (m.getHeaders().containsKey(MessageUtils.TARGET_PROTOCOL) && FunctionTypeUtils.isTypeCollection(this.outputType)) {
+					convertedOutput = this.convertMultipleOutputValuesIfNecessary(m.getPayload(), ObjectUtils.isEmpty(contentType) ? null : contentType);
+				}
+				else {
+					convertedOutput = this.convertOutputMessageIfNecessary(output, ObjectUtils.isEmpty(contentType) ? null : contentType[0]);
+				}
 			}
 			else if (output instanceof Collection && this.isOutputTypeMessage()) {
 				convertedOutput = this.convertMultipleOutputValuesIfNecessary(output, ObjectUtils.isEmpty(contentType) ? null : contentType);
diff --git a/spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistryTests.java b/spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistryTests.java
index 52105a174..0de0bb313 100644
--- a/spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistryTests.java
+++ b/spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistryTests.java
@@ -55,6 +55,7 @@ import org.springframework.cloud.function.context.FunctionCatalog;
 import org.springframework.cloud.function.context.FunctionRegistration;
 import org.springframework.cloud.function.context.FunctionRegistry;
 import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry.FunctionInvocationWrapper;
+import org.springframework.cloud.function.context.message.MessageUtils;
 import org.springframework.cloud.function.json.JsonMapper;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.annotation.Bean;
@@ -431,6 +432,18 @@ public class BeanFactoryAwareFunctionRegistryTests {
 		assertThat(result instanceof Message).isFalse();
 	}
 
+	@Test
+	public void testMultipleValuesInOutputHandlingWithTargetProtocol() throws Exception {
+		FunctionCatalog catalog = this.configureCatalog(CollectionOutConfiguration.class);
+		FunctionInvocationWrapper function = catalog.lookup("parseToList", "application/json");
+		assertThat(function).isNotNull();
+		Object result = function.apply(MessageBuilder.withPayload("1, 2, 3".getBytes())
+				.setHeader(MessageHeaders.CONTENT_TYPE, "text/plain")
+				.setHeader(MessageUtils.TARGET_PROTOCOL, "integration")
+				.build());
+		assertThat(result instanceof List).isTrue();
+	}
+
 	/**
 	 * The following two tests test the fallback mechanism when an accept header has several values.
 	 * The function produces Integer, which cannot be serialized by the default converter supporting text/plain