More header manipulation, limit bytes-per-second.

Add setHeader() and removeHeader() that mutate matching MockResponse
headers. Support "bytes-per-second" to throttle response rate. Split
request line into components.

Change-Id: I107baa3727ced2d7ddcc5e3cb983f5dc9341dabd

3 files changed, 68 insertions, 1 deletions

diff --git a/src/main/java/com/google/mockwebserver/MockResponse.java b/src/main/java/com/google/mockwebserver/MockResponse.java
index 5d06275..17e00f3 100644
--- a/src/main/java/com/google/mockwebserver/MockResponse.java
+++ b/src/main/java/com/google/mockwebserver/MockResponse.java
@@ -21,6 +21,7 @@ import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.UnsupportedEncodingException;
 import java.util.ArrayList;
+import java.util.Iterator;
 import java.util.List;
 
 /**
@@ -34,6 +35,7 @@ public final class MockResponse implements Cloneable {
     private String status = "HTTP/1.1 200 OK";
     private List<String> headers = new ArrayList<String>();
     private byte[] body = EMPTY_BODY;
+    private int bytesPerSecond = Integer.MAX_VALUE;
     private SocketPolicy socketPolicy = SocketPolicy.KEEP_OPEN;
 
     public MockResponse() {
@@ -84,6 +86,26 @@ public final class MockResponse implements Cloneable {
         return this;
     }
 
+    public MockResponse addHeader(String name, Object value) {
+        return addHeader(name + ": " + String.valueOf(value));
+    }
+
+    public MockResponse setHeader(String name, Object value) {
+        removeHeader(name);
+        return addHeader(name, value);
+    }
+
+    public MockResponse removeHeader(String name) {
+        name += ": ";
+        for (Iterator<String> i = headers.iterator(); i.hasNext();) {
+            String header = i.next();
+            if (name.regionMatches(true, 0, header, 0, name.length())) {
+                i.remove();
+            }
+        }
+        return this;
+    }
+
     /**
      * Returns an input stream containing the raw HTTP payload.
      */
@@ -140,6 +162,18 @@ public final class MockResponse implements Cloneable {
         return this;
     }
 
+    public int getBytesPerSecond() {
+        return bytesPerSecond;
+    }
+
+    /**
+     * Set simulated network speed, in bytes per second.
+     */
+    public MockResponse setBytesPerSecond(int bytesPerSecond) {
+        this.bytesPerSecond = bytesPerSecond;
+        return this;
+    }
+
     @Override public String toString() {
         return status;
     }
diff --git a/src/main/java/com/google/mockwebserver/MockWebServer.java b/src/main/java/com/google/mockwebserver/MockWebServer.java
index 65f4547..696a440 100644
--- a/src/main/java/com/google/mockwebserver/MockWebServer.java
+++ b/src/main/java/com/google/mockwebserver/MockWebServer.java
@@ -18,6 +18,7 @@ package com.google.mockwebserver;
 
 import static com.google.mockwebserver.SocketPolicy.DISCONNECT_AT_START;
 import static com.google.mockwebserver.SocketPolicy.FAIL_HANDSHAKE;
+
 import java.io.BufferedInputStream;
 import java.io.BufferedOutputStream;
 import java.io.ByteArrayOutputStream;
@@ -341,6 +342,7 @@ public final class MockWebServer {
                 } else if (response.getSocketPolicy() == SocketPolicy.SHUTDOWN_OUTPUT_AT_END) {
                     socket.shutdownOutput();
                 }
+                logger.info("Received request: " + request + " and responded: " + response);
                 sequenceNumber++;
                 return true;
             }
@@ -475,8 +477,24 @@ public final class MockWebServer {
             out.write((header + "\r\n").getBytes(ASCII));
         }
         out.write(("\r\n").getBytes(ASCII));
-        out.write(response.getBody());
         out.flush();
+
+        byte[] body = response.getBody();
+        int bytesPerSecond = response.getBytesPerSecond();
+
+        for (int offset = 0; offset < body.length; offset += bytesPerSecond) {
+            int count = Math.min(body.length - offset, bytesPerSecond);
+            out.write(body, offset, count);
+            out.flush();
+
+            if (offset + count < body.length) {
+                try {
+                    Thread.sleep(1000);
+                } catch (InterruptedException e) {
+                    throw new AssertionError();
+                }
+            }
+        }
     }
 
     /**
diff --git a/src/main/java/com/google/mockwebserver/RecordedRequest.java b/src/main/java/com/google/mockwebserver/RecordedRequest.java
index a06c0bc..8d7dc87 100644
--- a/src/main/java/com/google/mockwebserver/RecordedRequest.java
+++ b/src/main/java/com/google/mockwebserver/RecordedRequest.java
@@ -25,6 +25,8 @@ import javax.net.ssl.SSLSocket;
  */
 public final class RecordedRequest {
     private final String requestLine;
+    private final String method;
+    private final String path;
     private final List<String> headers;
     private final List<Integer> chunkSizes;
     private final int bodySize;
@@ -47,6 +49,11 @@ public final class RecordedRequest {
         } else {
             sslProtocol = null;
         }
+
+        int methodEnd = requestLine.indexOf(' ');
+        int pathEnd = requestLine.indexOf(' ', methodEnd + 1);
+        this.method = requestLine.substring(0, methodEnd);
+        this.path = requestLine.substring(methodEnd + 1, pathEnd);
     }
 
     public String getRequestLine() {
@@ -100,4 +107,12 @@ public final class RecordedRequest {
     @Override public String toString() {
         return requestLine;
     }
+
+    public String getMethod() {
+        return method;
+    }
+
+    public String getPath() {
+        return path;
+    }
 }