From 919b9902e9085c948800c1f830e7cb89f0ff1f00 Mon Sep 17 00:00:00 2001
From: Oleg Zhurakousky <ozhurakousky@pivotal.io>
Date: Wed, 18 Mar 2020 17:01:55 +0100
Subject: [PATCH] Improved handling of output Message conversion

remove the check for 'accept' header
added test

Resolves #465
---
 .../BeanFactoryAwareFunctionRegistry.java     | 18 +++++++-------
 ...BeanFactoryAwareFunctionRegistryTests.java | 24 ++++++++++++++++++-
 2 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistry.java b/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistry.java
index 9844e3242..38892a167 100644
--- a/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistry.java
+++ b/spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistry.java
@@ -735,16 +735,14 @@ public class BeanFactoryAwareFunctionRegistry
 			Message outputMessage = null;
 			if (value instanceof Message) {
 				MessageHeaders headers = ((Message) value).getHeaders();
-				if (!headers.containsKey("accept")) {
-					Map<String, Object> headersMap = (Map<String, Object>) ReflectionUtils
-						.getField(this.headersField, headers);
-					headersMap.put("accept", acceptedContentType);
-					// Set the contentType header to the value of accept for "legacy" reasons. But, do not set the
-					// contentType header to the value of accept if it is a wildcard type, as this doesn't make sense.
-					// This also applies to the else branch below.
-					if (acceptedContentType.isConcrete()) {
-						headersMap.put(MessageHeaders.CONTENT_TYPE, acceptedContentType);
-					}
+				Map<String, Object> headersMap = (Map<String, Object>) ReflectionUtils
+					.getField(this.headersField, headers);
+				headersMap.put("accept", acceptedContentType);
+				// Set the contentType header to the value of accept for "legacy" reasons. But, do not set the
+				// contentType header to the value of accept if it is a wildcard type, as this doesn't make sense.
+				// This also applies to the else branch below.
+				if (acceptedContentType.isConcrete()) {
+					headersMap.put(MessageHeaders.CONTENT_TYPE, acceptedContentType);
 				}
 			}
 			else {
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 1d7e630b7..f657db5af 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
@@ -364,7 +364,7 @@ public class BeanFactoryAwareFunctionRegistryTests {
 	}
 
 	/**
-	 * This test tests the fallback mechanism when an accept header has several values.
+	 * 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
 	 * (StringMessageConverter) but can by the one supporting application/json, which comes second.
 	 */
@@ -377,6 +377,16 @@ public class BeanFactoryAwareFunctionRegistryTests {
 		assertThat(result.getPayload()).isEqualTo("5".getBytes("UTF-8"));
 	}
 
+	@Test
+	public void testMultipleOrderedAcceptValuesMessageOutput() throws Exception {
+		FunctionCatalog catalog = this.configureCatalog(MultipleOrderedAcceptValuesAsMessageOutputConfiguration.class);
+		Function<String, Message<byte[]>> function = catalog.lookup(
+				"beanFactoryAwareFunctionRegistryTests.MultipleOrderedAcceptValuesAsMessageOutputConfiguration",
+				"text/plain,application/json");
+		assertThat(function).isNotNull();
+		Message<byte[]> result = function.apply("hello");
+		assertThat(result.getPayload()).isEqualTo("5".getBytes("UTF-8"));
+	}
 
 	@SuppressWarnings("unchecked")
 	@Test
@@ -775,4 +785,16 @@ public class BeanFactoryAwareFunctionRegistryTests {
 			return t.length();
 		}
 	}
+
+	@EnableAutoConfiguration
+	@Configuration
+	@Component
+	public static class MultipleOrderedAcceptValuesAsMessageOutputConfiguration implements Function<String, Message<Integer>> {
+
+		@Override
+		public Message<Integer> apply(String t) {
+			return MessageBuilder.withPayload(t.length()).build();
+		}
+
+	}
 }