From f0f654588087d8a21d7bb7074a72c046ded58984 Mon Sep 17 00:00:00 2001
From: David Costanzo <dcostanz@scharp.org>
Date: Wed, 22 Nov 2017 11:26:46 -0800
Subject: [PATCH] [#2140] Make FunctionalTest compatible with
 Response.writeChunk()

---
 framework/src/play/test/FunctionalTest.java   | 26 +++++++++++++++++++
 .../test/FunctionalTestTest.java              | 13 ++++++++++
 2 files changed, 39 insertions(+)

diff --git a/framework/src/play/test/FunctionalTest.java b/framework/src/play/test/FunctionalTest.java
index 36047e4efb..f4c144b4b4 100644
--- a/framework/src/play/test/FunctionalTest.java
+++ b/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;
     }
 
diff --git a/samples-and-tests/just-test-cases/test/FunctionalTestTest.java b/samples-and-tests/just-test-cases/test/FunctionalTestTest.java
index 3916be0379..64e87c85ab 100644
--- a/samples-and-tests/just-test-cases/test/FunctionalTestTest.java
+++ b/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);
+    }
 }