Merge pull request square#1286 from square/jw/response-body-source

Change MockResponse body to Source.

benchmarks/src/main/java/com/squareup/okhttp/benchmarks/Benchmark.java

@@ -24,9 +24,7 @@
 import com.squareup.okhttp.mockwebserver.MockResponse;
 import com.squareup.okhttp.mockwebserver.MockWebServer;
 import com.squareup.okhttp.mockwebserver.RecordedRequest;
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
-import java.io.OutputStream;
 import java.net.URL;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -35,9 +33,9 @@
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
 import java.util.logging.Logger;
-import java.util.zip.GZIPOutputStream;
 import javax.net.ssl.SSLContext;
 import okio.Buffer;
+import okio.GzipSink;
 
 /**
  * This benchmark is fake, but may be useful for certain relative comparisons.
@@ -180,18 +178,23 @@ private MockWebServer startServer() throws IOException {
   }
 
   private MockResponse newResponse() throws IOException {
-    byte[] body = new byte[bodyByteCount];
-    random.nextBytes(body);
+    byte[] bytes = new byte[bodyByteCount];
+    random.nextBytes(bytes);
+    Buffer body = new Buffer().write(bytes);
 
     MockResponse result = new MockResponse();
 
     if (gzip) {
-      body = gzip(body);
+      Buffer gzipBody = new Buffer();
+      GzipSink gzipSink = new GzipSink(gzipBody);
+      gzipSink.write(body, body.size());
+      gzipSink.close();
+      body = gzipBody;
       result.addHeader("Content-Encoding: gzip");
     }
 
     if (chunked) {
-      result.setChunkedBody(new Buffer().write(body), 1024);
+      result.setChunkedBody(body, 1024);
     } else {
       result.setBody(body);
     }
@@ -211,13 +214,4 @@ private String randomString(int length) {
     }
     return new String(result);
   }
-
-  /** Returns a gzipped copy of {@code bytes}. */
-  private byte[] gzip(byte[] bytes) throws IOException {
-    ByteArrayOutputStream bytesOut = new ByteArrayOutputStream();
-    OutputStream gzippedOut = new GZIPOutputStream(bytesOut);
-    gzippedOut.write(bytes);
-    gzippedOut.close();
-    return bytesOut.toByteArray();
-  }
 }

mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockResponse.java

@@ -16,14 +16,11 @@
 package com.squareup.okhttp.mockwebserver;
 
 import com.squareup.okhttp.internal.ws.WebSocketListener;
-import java.io.InputStream;
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 import okio.Buffer;
-import okio.Okio;
-import okio.Source;
 
 /** A scripted response to be replayed by the mock web server. */
 public final class MockResponse implements Cloneable {
@@ -32,10 +29,7 @@ public final class MockResponse implements Cloneable {
   private String status = "HTTP/1.1 200 OK";
   private List<String> headers = new ArrayList<>();
 
-  /** The response body content, or null if {@code bodyStream} is set. */
   private Buffer body;
-  /** The response body content, or null if {@code body} is set. */
-  private InputStream bodyStream;
 
   private int throttleBytesPerPeriod = Integer.MAX_VALUE;
   private long throttlePeriod = 1;
@@ -50,7 +44,7 @@ public final class MockResponse implements Cloneable {
 
   /** Creates a new mock response with an empty body. */
   public MockResponse() {
-    setBody(new Buffer());
+    setHeader("Content-Length", 0);
   }
 
   @Override public MockResponse clone() {
@@ -130,30 +124,14 @@ public MockResponse removeHeader(String name) {
     return this;
   }
 
-  /** Returns the raw HTTP payload, or null if this response is streamed. */
+  /** Returns a copy of the raw HTTP payload. */
   public Buffer getBody() {
-    return body != null ? body.clone() : null; // Defensive copy.
-  }
-
-  Source getBodySource() {
-    return bodyStream != null ? Okio.source(bodyStream) : getBody();
-  }
-
-  public MockResponse setBody(byte[] body) {
-    return setBody(new Buffer().write(body));
+    return body != null ? body.clone() : null;
   }
 
   public MockResponse setBody(Buffer body) {
     setHeader("Content-Length", body.size());
     this.body = body.clone(); // Defensive copy.
-    this.bodyStream = null;
-    return this;
-  }
-
-  public MockResponse setBody(InputStream bodyStream, long bodyLength) {
-    setHeader("Content-Length", bodyLength);
-    this.body = null;
-    this.bodyStream = bodyStream;
     return this;
   }
 
@@ -262,7 +240,6 @@ public MockResponse withWebSocketUpgrade(WebSocketListener listener) {
     setHeader("Connection", "Upgrade");
     setHeader("Upgrade", "websocket");
     body = null;
-    bodyStream = null;
     webSocketListener = listener;
     return this;
   }

mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockWebServer.java

@@ -73,7 +73,6 @@
 import okio.ByteString;
 import okio.Okio;
 import okio.Sink;
-import okio.Source;
 import okio.Timeout;
 
 import static com.squareup.okhttp.mockwebserver.SocketPolicy.DISCONNECT_AT_START;
@@ -668,10 +667,10 @@ private void writeHttpResponse(Socket socket, BufferedSink sink, MockResponse re
     sink.writeUtf8("\r\n");
     sink.flush();
 
-    Source source = response.getBodySource();
-    if (source == null) return;
+    Buffer body = response.getBody();
+    if (body == null) return;
     sleepIfDelayed(response);
-    throttledTransfer(response, socket, Okio.buffer(source), sink, Long.MAX_VALUE);
+    throttledTransfer(response, socket, body, sink, Long.MAX_VALUE);
   }
 
   private void sleepIfDelayed(MockResponse response) {
@@ -855,12 +854,12 @@ private void writeResponse(SpdyStream stream, MockResponse response) throws IOEx
         }
         spdyHeaders.add(new Header(headerParts[0], headerParts[1]));
       }
+
       Buffer body = response.getBody();
-      if (body == null) body = new Buffer();
-      boolean closeStreamAfterHeaders = body.size() > 0 || !response.getPushPromises().isEmpty();
+      boolean closeStreamAfterHeaders = body != null || !response.getPushPromises().isEmpty();
       stream.reply(spdyHeaders, closeStreamAfterHeaders);
       pushPromises(stream, response.getPushPromises());
-      if (body.size() > 0) {
+      if (body != null) {
         BufferedSink sink = Okio.buffer(stream.getSink());
         sleepIfDelayed(response);
         throttledTransfer(response, socket, body, sink, bodyLimit);
@@ -890,9 +889,9 @@ private void pushPromises(SpdyStream stream, List<PushPromise> promises) throws
         List<Integer> chunkSizes = Collections.emptyList(); // No chunked encoding for SPDY.
         requestQueue.add(new RecordedRequest(requestLine, pushPromise.getHeaders(), chunkSizes, 0,
             new Buffer(), sequenceNumber.getAndIncrement(), socket));
-        Buffer pushedBody = pushPromise.getResponse().getBody();
+        boolean hasBody = pushPromise.getResponse().getBody() != null;
         SpdyStream pushedStream =
-            stream.getConnection().pushStream(stream.getId(), pushedHeaders, pushedBody.size() > 0);
+            stream.getConnection().pushStream(stream.getId(), pushedHeaders, hasBody);
         writeResponse(pushedStream, pushPromise.getResponse());
       }
     }

mockwebserver/src/test/java/com/squareup/okhttp/mockwebserver/MockWebServerTest.java

@@ -17,7 +17,6 @@
 
 import com.squareup.okhttp.mockwebserver.rule.MockWebServerRule;
 import java.io.BufferedReader;
-import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
@@ -30,8 +29,6 @@
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 import okio.Buffer;
-import okio.BufferedSource;
-import okio.Okio;
 import org.junit.Rule;
 import org.junit.Test;
 
@@ -73,11 +70,7 @@ public final class MockWebServerTest {
   @Test public void setBodyAdjustsHeaders() throws IOException {
     MockResponse response = new MockResponse().setBody("ABC");
     assertEquals(Arrays.asList("Content-Length: 3"), response.getHeaders());
-    BufferedSource in = Okio.buffer(response.getBodySource());
-    assertEquals('A', in.readByte());
-    assertEquals('B', in.readByte());
-    assertEquals('C', in.readByte());
-    assertTrue(in.exhausted());
+    assertEquals("ABC", response.getBody().readUtf8());
     assertEquals("HTTP/1.1 200 OK", response.getStatus());
   }
 
@@ -210,18 +203,6 @@ public final class MockWebServerTest {
     server.getUrl("/b").openConnection().getInputStream(); // Should succeed.
   }
 
-  @Test public void streamingResponseBody() throws Exception {
-    InputStream responseBody = new ByteArrayInputStream("ABC".getBytes("UTF-8"));
-    server.enqueue(new MockResponse().setBody(responseBody, 3));
-
-    InputStream in = server.getUrl("/").openConnection().getInputStream();
-    assertEquals('A', in.read());
-    assertEquals('B', in.read());
-    assertEquals('C', in.read());
-
-    assertEquals(-1, responseBody.read()); // The body is exhausted.
-  }
-
   /**
    * Throttle the request body by sleeping 500ms after every 3 bytes. With a
    * 6-byte request, this should yield one sleep for a total delay of 500ms.

okhttp-tests/src/test/java/com/squareup/okhttp/internal/http/DisconnectTest.java

@@ -29,6 +29,7 @@
 import java.net.Socket;
 import java.util.concurrent.TimeUnit;
 
+import okio.Buffer;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -97,7 +98,7 @@ protected void configureSocket(Socket socket) throws IOException {
     int responseBodySize = 2 * 1024 * 1024; // 2 MiB
 
     server.enqueue(new MockResponse()
-        .setBody(new byte[responseBodySize])
+        .setBody(new Buffer().write(new byte[responseBodySize]))
         .throttleBody(64 * 1024, 125, TimeUnit.MILLISECONDS)); // 500 Kbps
     server.play();
 

okhttp-tests/src/test/java/com/squareup/okhttp/internal/http/ThreadInterruptTest.java

@@ -27,12 +27,11 @@
 import java.io.InterruptedIOException;
 import java.io.OutputStream;
 import java.net.HttpURLConnection;
-import java.net.InetAddress;
 import java.net.ServerSocket;
 import java.net.Socket;
-import java.net.UnknownHostException;
 import java.util.concurrent.TimeUnit;
 
+import okio.Buffer;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -101,7 +100,7 @@ protected void configureSocket(Socket socket) throws IOException {
     int responseBodySize = 2 * 1024 * 1024; // 2 MiB
 
     server.enqueue(new MockResponse()
-        .setBody(new byte[responseBodySize])
+        .setBody(new Buffer().write(new byte[responseBodySize]))
         .throttleBody(64 * 1024, 125, TimeUnit.MILLISECONDS)); // 500 Kbps
     server.play();
 

okhttp-tests/src/test/java/com/squareup/okhttp/internal/http/URLConnectionTest.java

@@ -741,9 +741,8 @@ public Socket createSocket(String host, int port, InetAddress localHost, int loc
   @Test public void contentDisagreesWithChunkedHeader() throws IOException {
     MockResponse mockResponse = new MockResponse();
     mockResponse.setChunkedBody("abc", 3);
-    Buffer buffer = new Buffer();
-    buffer.writeAll(mockResponse.getBody());
-    buffer.write("\r\nYOU SHOULD NOT SEE THIS".getBytes("UTF-8"));
+    Buffer buffer = mockResponse.getBody();
+    buffer.writeUtf8("\r\nYOU SHOULD NOT SEE THIS");
     mockResponse.setBody(buffer);
     mockResponse.clearHeaders();
     mockResponse.addHeader("Transfer-encoding: chunked");