Expose Headers directly on MockResponse.

wuqingman · Jan 5, 2015 · c5c4e78 · c5c4e78
1 parent 6763803
commit c5c4e78
Show file tree

Hide file tree

Showing 7 changed files with 32 additions and 61 deletions.
diff --git a/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockResponse.java b/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockResponse.java
@@ -73,17 +73,7 @@ public MockResponse setStatus(String status) {
   }
 
   /** Returns the HTTP headers, such as "Content-Length: 0". */
-  public List<String> getHeaders() {
-    Headers headers = this.headers.build();
-    int size = headers.size();
-    List<String> headerList = new ArrayList<>(size);
-    for (int i = 0; i < size; i++) {
-      headerList.add(headers.name(i) + ": " + headers.value(i));
-    }
-    return headerList;
-  }
-
-  Headers getNewHeaders() {
+  public Headers getHeaders() {
     return headers.build();
   }
 

diff --git a/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockWebServer.java b/mockwebserver/src/main/java/com/squareup/okhttp/mockwebserver/MockWebServer.java
@@ -611,7 +611,7 @@ private void handleWebSocketUpgrade(Socket socket, BufferedSource source, Buffer
     final Response fancyResponse = new Response.Builder()
         .code(Integer.parseInt(response.getStatus().split(" ")[1]))
         .message(response.getStatus().split(" ", 3)[2])
-        .headers(response.getNewHeaders())
+        .headers(response.getHeaders())
         .request(fancyRequest)
         .protocol(Protocol.HTTP_1_1)
         .build();
@@ -648,9 +648,11 @@ private void writeHttpResponse(Socket socket, BufferedSink sink, MockResponse re
     sink.writeUtf8(response.getStatus());
     sink.writeUtf8("\r\n");
 
-    List<String> headers = response.getHeaders();
+    Headers headers = response.getHeaders();
     for (int i = 0, size = headers.size(); i < size; i++) {
-      sink.writeUtf8(headers.get(i));
+      sink.writeUtf8(headers.name(i));
+      sink.writeUtf8(": ");
+      sink.writeUtf8(headers.value(i));
       sink.writeUtf8("\r\n");
     }
     sink.writeUtf8("\r\n");
@@ -834,7 +836,7 @@ private void writeResponse(SpdyStream stream, MockResponse response) throws IOEx
       if (protocol == Protocol.SPDY_3) {
         spdyHeaders.add(new Header(Header.VERSION, statusParts[0]));
       }
-      Headers headers = response.getNewHeaders();
+      Headers headers = response.getHeaders();
       for (int i = 0, size = headers.size(); i < size; i++) {
         spdyHeaders.add(new Header(headers.name(i), headers.value(i)));
       }

diff --git a/mockwebserver/src/test/java/com/squareup/okhttp/mockwebserver/MockWebServerTest.java b/mockwebserver/src/test/java/com/squareup/okhttp/mockwebserver/MockWebServerTest.java
@@ -25,6 +25,7 @@
 import java.net.SocketTimeoutException;
 import java.net.URL;
 import java.net.URLConnection;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
@@ -64,13 +65,13 @@ public final class MockWebServerTest {
 
   @Test public void defaultMockResponse() {
     MockResponse response = new MockResponse();
-    assertEquals(Arrays.asList("Content-Length: 0"), response.getHeaders());
+    assertEquals(Arrays.asList("Content-Length: 0"), headersToList(response));
     assertEquals("HTTP/1.1 200 OK", response.getStatus());
   }
 
   @Test public void setBodyAdjustsHeaders() throws IOException {
     MockResponse response = new MockResponse().setBody("ABC");
-    assertEquals(Arrays.asList("Content-Length: 3"), response.getHeaders());
+    assertEquals(Arrays.asList("Content-Length: 3"), headersToList(response));
     assertEquals("ABC", response.getBody().readUtf8());
     assertEquals("HTTP/1.1 200 OK", response.getStatus());
   }
@@ -80,7 +81,7 @@ public final class MockWebServerTest {
         .clearHeaders()
         .addHeader("Cookie: s=square")
         .addHeader("Cookie", "a=android");
-    assertEquals(Arrays.asList("Cookie: s=square", "Cookie: a=android"), response.getHeaders());
+    assertEquals(Arrays.asList("Cookie: s=square", "Cookie: a=android"), headersToList(response));
   }
 
   @Test public void mockResponseSetHeader() {
@@ -90,7 +91,7 @@ public final class MockWebServerTest {
         .addHeader("Cookie: a=android")
         .addHeader("Cookies: delicious");
     response.setHeader("cookie", "r=robot");
-    assertEquals(Arrays.asList("Cookies: delicious", "cookie: r=robot"), response.getHeaders());
+    assertEquals(Arrays.asList("Cookies: delicious", "cookie: r=robot"), headersToList(response));
   }
 
   @Test public void regularResponse() throws Exception {
@@ -275,4 +276,14 @@ public final class MockWebServerTest {
     assertTrue(String.format("Request + Response: %sms", elapsedMillis), elapsedMillis >= 1000);
     assertTrue(String.format("Request + Response: %sms", elapsedMillis), elapsedMillis <= 1100);
   }
+
+  private List<String> headersToList(MockResponse response) {
+    Headers headers = response.getHeaders();
+    int size = headers.size();
+    List<String> headerList = new ArrayList<>(size);
+    for (int i = 0; i < size; i++) {
+      headerList.add(headers.name(i) + ": " + headers.value(i));
+    }
+    return headerList;
+  }
 }
diff --git a/okhttp-tests/src/test/java/com/squareup/okhttp/CacheTest.java b/okhttp-tests/src/test/java/com/squareup/okhttp/CacheTest.java
@@ -2214,14 +2214,11 @@ private void assertFullyCached(MockResponse response) throws Exception {
    */
   private MockResponse truncateViolently(MockResponse response, int numBytesToKeep) {
     response.setSocketPolicy(DISCONNECT_AT_END);
-    List<String> headers = new ArrayList<>(response.getHeaders());
+    Headers headers = response.getHeaders();
     Buffer truncatedBody = new Buffer();
     truncatedBody.write(response.getBody(), numBytesToKeep);
     response.setBody(truncatedBody);
-    response.clearHeaders();
-    for (String header : headers) {
-      response.addHeader(header);
-    }
+    response.setHeaders(headers);
     return response;
   }
 
@@ -2241,12 +2238,7 @@ enum TransferKind {
       @Override void setBody(MockResponse response, Buffer content, int chunkSize) {
         response.setBody(content);
         response.setSocketPolicy(DISCONNECT_AT_END);
-        for (Iterator<String> h = response.getHeaders().iterator(); h.hasNext(); ) {
-          if (h.next().startsWith("Content-Length:")) {
-            h.remove();
-            break;
-          }
-        }
+        response.removeHeader("Content-Length");
       }
     };
 

diff --git a/okhttp-tests/src/test/java/com/squareup/okhttp/internal/http/URLConnectionTest.java b/okhttp-tests/src/test/java/com/squareup/okhttp/internal/http/URLConnectionTest.java
@@ -68,7 +68,6 @@
 import java.util.Collections;
 import java.util.EnumSet;
 import java.util.HashSet;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Random;
@@ -3117,12 +3116,7 @@ enum TransferKind {
       @Override void setBody(MockResponse response, Buffer content, int chunkSize) {
         response.setBody(content);
         response.setSocketPolicy(DISCONNECT_AT_END);
-        for (Iterator<String> h = response.getHeaders().iterator(); h.hasNext(); ) {
-          if (h.next().startsWith("Content-Length:")) {
-            h.remove();
-            break;
-          }
-        }
+        response.removeHeader("Content-Length");
       }
       @Override void setForRequest(HttpURLConnection connection, int contentLength) {
       }

diff --git a/okhttp-urlconnection/src/test/java/com/squareup/okhttp/UrlConnectionCacheTest.java b/okhttp-urlconnection/src/test/java/com/squareup/okhttp/UrlConnectionCacheTest.java
@@ -44,7 +44,6 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Date;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Locale;
 import java.util.TimeZone;
@@ -1744,14 +1743,11 @@ private void assertFullyCached(MockResponse response) throws Exception {
    */
   private MockResponse truncateViolently(MockResponse response, int numBytesToKeep) {
     response.setSocketPolicy(DISCONNECT_AT_END);
-    List<String> headers = new ArrayList<>(response.getHeaders());
+    Headers headers = response.getHeaders();
     Buffer truncatedBody = new Buffer();
     truncatedBody.write(response.getBody(), numBytesToKeep);
     response.setBody(truncatedBody);
-    response.clearHeaders();
-    for (String header : headers) {
-      response.addHeader(header);
-    }
+    response.setHeaders(headers);
     return response;
   }
 
@@ -1813,12 +1809,7 @@ enum TransferKind {
       @Override void setBody(MockResponse response, Buffer content, int chunkSize) {
         response.setBody(content);
         response.setSocketPolicy(DISCONNECT_AT_END);
-        for (Iterator<String> h = response.getHeaders().iterator(); h.hasNext(); ) {
-          if (h.next().startsWith("Content-Length:")) {
-            h.remove();
-            break;
-          }
-        }
+        response.removeHeader("Content-Length");
       }
     };
 

diff --git a/okhttp-urlconnection/src/test/java/com/squareup/okhttp/internal/huc/ResponseCacheTest.java b/okhttp-urlconnection/src/test/java/com/squareup/okhttp/internal/huc/ResponseCacheTest.java
@@ -52,7 +52,6 @@
 import java.util.Arrays;
 import java.util.Date;
 import java.util.HashMap;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
@@ -1365,14 +1364,11 @@ private void assertFullyCached(MockResponse response) throws Exception {
    */
   private MockResponse truncateViolently(MockResponse response, int numBytesToKeep) {
     response.setSocketPolicy(DISCONNECT_AT_END);
-    List<String> headers = new ArrayList<>(response.getHeaders());
+    Headers headers = response.getHeaders();
     Buffer truncatedBody = new Buffer();
     truncatedBody.write(response.getBody(), numBytesToKeep);
     response.setBody(truncatedBody);
-    response.clearHeaders();
-    for (String header : headers) {
-      response.addHeader(header);
-    }
+    response.setHeaders(headers);
     return response;
   }
 
@@ -1433,12 +1429,7 @@ enum TransferKind {
       @Override void setBody(MockResponse response, Buffer content, int chunkSize) {
         response.setBody(content);
         response.setSocketPolicy(DISCONNECT_AT_END);
-        for (Iterator<String> h = response.getHeaders().iterator(); h.hasNext(); ) {
-          if (h.next().startsWith("Content-Length:")) {
-            h.remove();
-            break;
-          }
-        }
+        response.removeHeader("Content-Length");
       }
     };