From 425be34b985fa6bea94001f4abd3e9a920a546a3 Mon Sep 17 00:00:00 2001
From: Dave Syer <dsyer@pivotal.io>
Date: Mon, 23 Jan 2017 09:37:02 +0000
Subject: [PATCH] Send content type to web response if available

---
 .../function/web/flux/FluxResponseBodyEmitter.java | 14 ++++++++++++++
 .../cloud/function/web/RestApplicationTests.java   |  9 +++++----
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/spring-cloud-function-web/src/main/java/org/springframework/cloud/function/web/flux/FluxResponseBodyEmitter.java b/spring-cloud-function-web/src/main/java/org/springframework/cloud/function/web/flux/FluxResponseBodyEmitter.java
index ae449caa3..0f2f95ae2 100644
--- a/spring-cloud-function-web/src/main/java/org/springframework/cloud/function/web/flux/FluxResponseBodyEmitter.java
+++ b/spring-cloud-function-web/src/main/java/org/springframework/cloud/function/web/flux/FluxResponseBodyEmitter.java
@@ -16,7 +16,9 @@
 
 package org.springframework.cloud.function.web.flux;
 
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.MediaType;
+import org.springframework.http.server.ServerHttpResponse;
 import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;
 
 import reactor.core.publisher.Flux;
@@ -28,6 +30,8 @@ import reactor.core.publisher.Flux;
  */
 class FluxResponseBodyEmitter<T> extends ResponseBodyEmitter {
 
+	private MediaType mediaType;
+
 	public FluxResponseBodyEmitter(Flux<T> observable) {
 		this(null, null, observable);
 	}
@@ -35,7 +39,17 @@ class FluxResponseBodyEmitter<T> extends ResponseBodyEmitter {
 	public FluxResponseBodyEmitter(Long timeout, MediaType mediaType,
 			Flux<T> observable) {
 		super(timeout);
+		this.mediaType = mediaType;
 		new ResponseBodyEmitterSubscriber<>(mediaType, observable, this);
 	}
 
+	@Override
+	protected void extendResponse(ServerHttpResponse outputMessage) {
+		super.extendResponse(outputMessage);
+
+		HttpHeaders headers = outputMessage.getHeaders();
+		if (headers.getContentType() == null && this.mediaType!=null && !MediaType.ALL.equals(this.mediaType)) {
+			headers.setContentType(this.mediaType);
+		}
+	}
 }
diff --git a/spring-cloud-function-web/src/test/java/org/springframework/cloud/function/web/RestApplicationTests.java b/spring-cloud-function-web/src/test/java/org/springframework/cloud/function/web/RestApplicationTests.java
index 167c59a67..fb3b4369d 100644
--- a/spring-cloud-function-web/src/test/java/org/springframework/cloud/function/web/RestApplicationTests.java
+++ b/spring-cloud-function-web/src/test/java/org/springframework/cloud/function/web/RestApplicationTests.java
@@ -104,17 +104,18 @@ public class RestApplicationTests {
 						.accept(MediaType.APPLICATION_JSON).build(),
 				String.class);
 		assertThat(result.getBody()).isEqualTo("[[\"go\",\"home\"],[\"come\",\"back\"]]");
-		// TODO: test this
-		// assertThat(result.getHeaders().getContentType()).isEqualTo(MediaType.APPLICATION_JSON);
+		assertThat(result.getHeaders().getContentType()).isEqualTo(MediaType.APPLICATION_JSON);
 	}
 
 	@Test
 	public void sentencesAcceptSse() throws Exception {
-		assertThat(rest.exchange(
+		ResponseEntity<String> result = rest.exchange(
 				RequestEntity.get(new URI("http://localhost:" + port + "/sentences"))
 						.accept(EVENT_STREAM).build(),
-				String.class).getBody())
+				String.class);
+		assertThat(result.getBody())
 						.isEqualTo(sse("[\"go\",\"home\"]", "[\"come\",\"back\"]"));
+		assertThat(result.getHeaders().getContentType().isCompatibleWith(EVENT_STREAM)).isTrue();
 	}
 
 	@Test