[#2140] Make FunctionalTest compatible with Response.writeChunk()

framework/src/play/test/FunctionalTest.java

@@ -31,6 +31,7 @@
 import play.Invoker;
 import play.Invoker.InvocationContext;
 import play.classloading.enhancers.ControllersEnhancer.ControllerInstrumentation;
+import play.libs.F.Action;
 import play.mvc.ActionInvoker;
 import play.mvc.Controller;
 import play.mvc.Http;
@@ -399,6 +400,31 @@ public static Response makeRequest(Request request) {
     public static Response newResponse() {
         Response response = new Response();
         response.out = new ByteArrayOutputStream();
+
+        // Register an onWriteChunk action so that Response.writeChunk() won't throw
+        // an unhandled exception if the controller action calls it.
+        response.onWriteChunk(
+            new Action<Object>() {
+                @Override
+                public void invoke(Object chunk) {
+                    // Mimic the behavior of PlayHandler$LazyChunkedInput.writeChunk()
+                    if (chunk != null) {
+                        try {
+                            byte[] bytes;
+                            if (chunk instanceof byte[]) {
+                                bytes = (byte[]) chunk;
+                            } else {
+                                bytes = chunk.toString().getBytes(response.encoding);
+                            }
+                            response.out.write(bytes);
+                        } catch (Exception exception) {
+                            // Something is wrong with the chunk.
+                            throw new RuntimeException(exception);
+                        }
+                    }
+                }
+            });
+
         return response;
     }
 

samples-and-tests/just-test-cases/test/FunctionalTestTest.java

@@ -126,5 +126,18 @@ public void testGettingStaticFile() {
 	Response response = GET("http://localhost:9003/public/session.test?req=1");
 	assertIsOk(response);
     }
+
+    /**
+     * This is a regression test for [#2140], which is a bug in FunctionalTest that prevented it from
+     * testing a controller action that uses {@link Response#writeChunk(Object)}.
+     */
+    @Test
+    public void testWriteChunks() {
+        Response response = GET("/application/writeChunks");
+        assertTrue(response.chunked);
+        assertIsOk(response);
+        assertContentType("text/plain", response);
+        assertContentEquals("abcæøåæøå", response);
+    }
 }