3 files changed, 59 insertions, 23 deletions

diff --git a/tests/src/tests/http/MockResponse.java b/tests/src/tests/http/MockResponse.java
index 9893e2fc..21397019 100644
--- a/tests/src/tests/http/MockResponse.java
+++ b/tests/src/tests/http/MockResponse.java
@@ -16,28 +16,29 @@
 
 package tests.http;
 
+import static tests.http.MockWebServer.ASCII;
+
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.UnsupportedEncodingException;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
-
-import static tests.http.MockWebServer.ASCII;
+import java.util.Map;
 
 /**
  * A scripted response to be replayed by the mock web server.
  */
 public class MockResponse {
-    private static final String EMPTY_BODY_HEADER = "Content-Length: 0";
-    private static final String CHUNKED_BODY_HEADER = "Transfer-encoding: chunked";
     private static final byte[] EMPTY_BODY = new byte[0];
 
     private String status = "HTTP/1.1 200 OK";
-    private List<String> headers = new ArrayList<String>();
+    private Map<String, String> headers = new HashMap<String, String>();
     private byte[] body = EMPTY_BODY;
+    private boolean closeConnectionAfter = false;
 
     public MockResponse() {
-        headers.add(EMPTY_BODY_HEADER);
+        addHeader("Content-Length", 0);
     }
 
     /**
@@ -56,14 +57,22 @@ public class MockResponse {
      * Returns the HTTP headers, such as "Content-Length: 0".
      */
     public List<String> getHeaders() {
-        return headers;
+        List<String> headerStrings = new ArrayList<String>();
+        for (String header : headers.keySet()) {
+            headerStrings.add(header + ": " + headers.get(header));
+        }
+        return headerStrings;
     }
 
-    public MockResponse addHeader(String header) {
-        headers.add(header);
+    public MockResponse addHeader(String header, String value) {
+        headers.put(header.toLowerCase(), value);
         return this;
     }
 
+    public MockResponse addHeader(String header, int value) {
+        return addHeader(header, Integer.toString(value));
+    }
+
     /**
      * Returns an input stream containing the raw HTTP payload.
      */
@@ -72,10 +81,7 @@ public class MockResponse {
     }
 
     public MockResponse setBody(byte[] body) {
-        if (this.body == EMPTY_BODY) {
-            headers.remove(EMPTY_BODY_HEADER);
-        }
-        this.headers.add("Content-Length: " + body.length);
+        addHeader("Content-Length", body.length);
         this.body = body;
         return this;
     }
@@ -89,8 +95,7 @@ public class MockResponse {
     }
 
     public MockResponse setChunkedBody(byte[] body, int maxChunkSize) throws IOException {
-        headers.remove(EMPTY_BODY_HEADER);
-        headers.add(CHUNKED_BODY_HEADER);
+        addHeader("Transfer-encoding", "chunked");
 
         ByteArrayOutputStream bytesOut = new ByteArrayOutputStream();
         int pos = 0;
@@ -114,4 +119,13 @@ public class MockResponse {
     @Override public String toString() {
         return status;
     }
+
+    public boolean shouldCloseConnectionAfter() {
+        return closeConnectionAfter;
+    }
+
+    public MockResponse setCloseConnectionAfter(boolean closeConnectionAfter) {
+        this.closeConnectionAfter = closeConnectionAfter;
+        return this;
+    }
 }
diff --git a/tests/src/tests/http/MockWebServer.java b/tests/src/tests/http/MockWebServer.java
index e3df2e83..b2cb8d7b 100644
--- a/tests/src/tests/http/MockWebServer.java
+++ b/tests/src/tests/http/MockWebServer.java
@@ -32,21 +32,22 @@ import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-import java.util.concurrent.LinkedBlockingDeque;
 import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
 
 /**
  * A scriptable web server. Callers supply canned responses and the server
  * replays them upon request in sequence.
+ *
+ * TODO: merge with the version from libcore/support/src/tests/java once it's in.
  */
 public final class MockWebServer {
-
     static final String ASCII = "US-ASCII";
 
     private final BlockingQueue<RecordedRequest> requestQueue
             = new LinkedBlockingQueue<RecordedRequest>();
     private final BlockingQueue<MockResponse> responseQueue
-            = new LinkedBlockingDeque<MockResponse>();
+            = new LinkedBlockingQueue<MockResponse>();
     private int bodyLimit = Integer.MAX_VALUE;
     private final ExecutorService executor = Executors.newCachedThreadPool();
 
@@ -88,6 +89,16 @@ public final class MockWebServer {
         return requestQueue.take();
     }
 
+    public RecordedRequest takeRequestWithTimeout(long timeoutMillis) throws InterruptedException {
+        return requestQueue.poll(timeoutMillis, TimeUnit.MILLISECONDS);
+    }
+
+    public List<RecordedRequest> drainRequests() {
+        List<RecordedRequest> requests = new ArrayList<RecordedRequest>();
+        requestQueue.drainTo(requests);
+        return requests;
+    }
+
     /**
      * Starts the server, serves all enqueued requests, and shuts the server
      * down.
@@ -130,7 +141,11 @@ public final class MockWebServer {
                         }
                     }
                     requestQueue.add(request);
-                    writeResponse(out, computeResponse(request));
+                    MockResponse response = computeResponse(request);
+                    writeResponse(out, response);
+                    if (response.shouldCloseConnectionAfter()) {
+                        break;
+                    }
                     sequenceNumber++;
                 }
 
@@ -146,7 +161,7 @@ public final class MockWebServer {
      */
     private RecordedRequest readRequest(InputStream in, int sequenceNumber) throws IOException {
         String request = readAsciiUntilCrlf(in);
-        if (request.isEmpty()) {
+        if (request.equals("")) {
             return null; // end of data; no more requests
         }
 
@@ -154,7 +169,7 @@ public final class MockWebServer {
         int contentLength = -1;
         boolean chunked = false;
         String header;
-        while (!(header = readAsciiUntilCrlf(in)).isEmpty()) {
+        while (!(header = readAsciiUntilCrlf(in)).equals("")) {
             headers.add(header);
             String lowercaseHeader = header.toLowerCase();
             if (contentLength == -1 && lowercaseHeader.startsWith("content-length:")) {
@@ -219,7 +234,6 @@ public final class MockWebServer {
         }
         out.write(("\r\n").getBytes(ASCII));
         out.write(response.getBody());
-        out.write(("\r\n").getBytes(ASCII));
         out.flush();
     }
 
@@ -260,7 +274,7 @@ public final class MockWebServer {
 
     private void readEmptyLine(InputStream in) throws IOException {
         String line = readAsciiUntilCrlf(in);
-        if (!line.isEmpty()) {
+        if (!line.equals("")) {
             throw new IllegalStateException("Expected empty but was: " + line);
         }
     }
diff --git a/tests/src/tests/http/RecordedRequest.java b/tests/src/tests/http/RecordedRequest.java
index c8050065..6b67af20 100644
--- a/tests/src/tests/http/RecordedRequest.java
+++ b/tests/src/tests/http/RecordedRequest.java
@@ -82,4 +82,12 @@ public final class RecordedRequest {
     @Override public String toString() {
         return requestLine;
     }
+
+    public String getMethod() {
+        return getRequestLine().split(" ")[0];
+    }
+
+    public String getPath() {
+        return getRequestLine().split(" ")[1];
+    }
 }