Support StreamingResponseBody return value type

Issue: SPR-12831

spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/ServletInvocableHandlerMethodTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2014 the original author or authors.
+ * Copyright 2002-2015 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,10 +16,14 @@
 
 package org.springframework.web.servlet.mvc.method.annotation;
 
+import static org.junit.Assert.*;
+import static org.mockito.Mockito.*;
+
 import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+
 import javax.servlet.http.HttpServletResponse;
 
 import org.junit.Before;
@@ -46,8 +50,6 @@ import org.springframework.web.method.support.HandlerMethodReturnValueHandlerCom
 import org.springframework.web.method.support.ModelAndViewContainer;
 import org.springframework.web.servlet.view.RedirectView;
 
-import static org.junit.Assert.*;
-import static org.mockito.Mockito.*;
 
 /**
  * Test fixture with {@link ServletInvocableHandlerMethod}.
@@ -220,6 +222,30 @@ public class ServletInvocableHandlerMethodTests {
 		assertEquals("", this.response.getContentAsString());
 	}
 
+	@Test
+	public void wrapConcurrentResult_ResponseBodyEmitter() throws Exception {
+		List<HttpMessageConverter<?>> converters = new ArrayList<HttpMessageConverter<?>>();
+		converters.add(new StringHttpMessageConverter());
+		this.returnValueHandlers.addHandler(new ResponseBodyEmitterReturnValueHandler(converters));
+		ServletInvocableHandlerMethod handlerMethod = getHandlerMethod(new AsyncHandler(), "handleWithEmitter");
+		handlerMethod = handlerMethod.wrapConcurrentResult(null);
+		handlerMethod.invokeAndHandle(this.webRequest, this.mavContainer);
+
+		assertEquals(200, this.response.getStatus());
+		assertEquals("", this.response.getContentAsString());
+	}
+
+	@Test
+	public void wrapConcurrentResult_StreamingResponseBody() throws Exception {
+		this.returnValueHandlers.addHandler(new StreamingResponseBodyReturnValueHandler());
+		ServletInvocableHandlerMethod handlerMethod = getHandlerMethod(new AsyncHandler(), "handleWithStreaming");
+		handlerMethod = handlerMethod.wrapConcurrentResult(null);
+		handlerMethod.invokeAndHandle(this.webRequest, this.mavContainer);
+
+		assertEquals(200, this.response.getStatus());
+		assertEquals("", this.response.getContentAsString());
+	}
+
 	// SPR-12287 (16/Oct/14 comments)
 
 	@Test
@@ -273,6 +299,7 @@ public class ServletInvocableHandlerMethodTests {
 		}
 	}
 
+	@SuppressWarnings("unused")
 	private static class MethodLevelResponseBodyHandler {
 
 		@ResponseBody
@@ -281,28 +308,28 @@ public class ServletInvocableHandlerMethodTests {
 		}
 	}
 
+	@SuppressWarnings("unused")
 	@ResponseBody
 	private static class TypeLevelResponseBodyHandler {
 
-		@SuppressWarnings("unused")
 		public DeferredResult<String> handle() {
 			return new DeferredResult<String>();
 		}
 	}
 
+	@SuppressWarnings("unused")
 	private static class ResponseEntityHandler {
 
-		@SuppressWarnings("unused")
 		public DeferredResult<ResponseEntity<String>> handleDeferred() {
 			return new DeferredResult<>();
 		}
 
-		@SuppressWarnings("unused")
 		public ResponseEntity handleRawType() {
 			return ResponseEntity.ok().build();
 		}
 	}
 
+	@SuppressWarnings("unused")
 	private static class ExceptionRaisingReturnValueHandler implements HandlerMethodReturnValueHandler {
 
 		@Override
@@ -317,4 +344,16 @@ public class ServletInvocableHandlerMethodTests {
 		}
 	}
 
+	@SuppressWarnings("unused")
+	private static class AsyncHandler {
+
+		public ResponseBodyEmitter handleWithEmitter() {
+			return null;
+		}
+
+		public StreamingResponseBody handleWithStreaming() {
+			return null;
+		}
+	}
+
 }

spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/StreamingResponseBodyReturnValueHandlerTests.java

@@ -0,0 +1,170 @@
+/*
+ * Copyright 2002-2015 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.web.servlet.mvc.method.annotation;
+
+import static org.junit.Assert.*;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.lang.reflect.Method;
+import java.nio.charset.Charset;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import org.springframework.core.MethodParameter;
+import org.springframework.http.ResponseEntity;
+import org.springframework.mock.web.test.MockHttpServletRequest;
+import org.springframework.mock.web.test.MockHttpServletResponse;
+import org.springframework.web.context.request.NativeWebRequest;
+import org.springframework.web.context.request.ServletWebRequest;
+import org.springframework.web.context.request.async.AsyncWebRequest;
+import org.springframework.web.context.request.async.StandardServletAsyncWebRequest;
+import org.springframework.web.context.request.async.WebAsyncUtils;
+import org.springframework.web.method.support.ModelAndViewContainer;
+
+
+/**
+ * Unit tests for
+ * {@link org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBodyReturnValueHandler}.
+ *
+ * @author Rossen Stoyanchev
+ */
+public class StreamingResponseBodyReturnValueHandlerTests {
+
+	private StreamingResponseBodyReturnValueHandler handler;
+
+	private ModelAndViewContainer mavContainer;
+
+	private NativeWebRequest webRequest;
+
+	private MockHttpServletRequest request;
+
+	private MockHttpServletResponse response;
+
+
+	@Before
+	public void setUp() throws Exception {
+
+		this.handler = new StreamingResponseBodyReturnValueHandler();
+		this.mavContainer = new ModelAndViewContainer();
+
+		this.request = new MockHttpServletRequest("GET", "/path");
+		this.response = new MockHttpServletResponse();
+		this.webRequest = new ServletWebRequest(this.request, this.response);
+
+		AsyncWebRequest asyncWebRequest = new StandardServletAsyncWebRequest(this.request, this.response);
+		WebAsyncUtils.getAsyncManager(this.webRequest).setAsyncWebRequest(asyncWebRequest);
+		this.request.setAsyncSupported(true);
+	}
+
+	@Test
+	public void supportsReturnType() throws Exception {
+		assertTrue(this.handler.supportsReturnType(returnType(TestController.class, "handle")));
+		assertTrue(this.handler.supportsReturnType(returnType(TestController.class, "handleResponseEntity")));
+		assertFalse(this.handler.supportsReturnType(returnType(TestController.class, "handleResponseEntityString")));
+		assertFalse(this.handler.supportsReturnType(returnType(TestController.class, "handleResponseEntityParameterized")));
+	}
+
+	@Test
+	public void streamingResponseBody() throws Exception {
+
+		CountDownLatch latch = new CountDownLatch(1);
+
+		MethodParameter returnType = returnType(TestController.class, "handle");
+		StreamingResponseBody streamingBody = new StreamingResponseBody() {
+
+			@Override
+			public void writeTo(OutputStream outputStream) throws IOException {
+				outputStream.write("foo".getBytes(Charset.forName("UTF-8")));
+				latch.countDown();
+			}
+		};
+		this.handler.handleReturnValue(streamingBody, returnType, this.mavContainer, this.webRequest);
+
+		assertTrue(this.request.isAsyncStarted());
+		assertTrue(latch.await(5, TimeUnit.SECONDS));
+		assertEquals("foo", this.response.getContentAsString());
+	}
+
+
+	@Test
+	public void responseEntity() throws Exception {
+
+		CountDownLatch latch = new CountDownLatch(1);
+
+		MethodParameter returnType = returnType(TestController.class, "handleResponseEntity");
+		ResponseEntity<StreamingResponseBody> emitter = ResponseEntity.ok().header("foo", "bar")
+				.body(new StreamingResponseBody() {
+
+					@Override
+					public void writeTo(OutputStream outputStream) throws IOException {
+						outputStream.write("foo".getBytes(Charset.forName("UTF-8")));
+						latch.countDown();
+					}
+				});
+		this.handler.handleReturnValue(emitter, returnType, this.mavContainer, this.webRequest);
+
+		assertTrue(this.request.isAsyncStarted());
+		assertEquals(200, this.response.getStatus());
+		assertEquals("bar", this.response.getHeader("foo"));
+
+		assertTrue(latch.await(5, TimeUnit.SECONDS));
+		assertEquals("foo", this.response.getContentAsString());
+
+	}
+
+	@Test
+	public void responseEntityNoContent() throws Exception {
+		MethodParameter returnType = returnType(TestController.class, "handleResponseEntity");
+		ResponseEntity<?> emitter = ResponseEntity.noContent().build();
+		this.handler.handleReturnValue(emitter, returnType, this.mavContainer, this.webRequest);
+
+		assertFalse(this.request.isAsyncStarted());
+		assertEquals(204, this.response.getStatus());
+	}
+
+
+	private MethodParameter returnType(Class<?> clazz, String methodName) throws NoSuchMethodException {
+		Method method = clazz.getDeclaredMethod(methodName);
+		return new MethodParameter(method, -1);
+	}
+
+
+	@SuppressWarnings("unused")
+	private static class TestController {
+
+		private StreamingResponseBody handle() {
+			return null;
+		}
+
+		private ResponseEntity<StreamingResponseBody> handleResponseEntity() {
+			return null;
+		}
+
+		private ResponseEntity<String> handleResponseEntityString() {
+			return null;
+		}
+
+		private ResponseEntity<AtomicReference<String>> handleResponseEntityParameterized() {
+			return null;
+		}
+	}
+
+}