Home | History | Annotate | Download | only in mockwebserver
      1 /*
      2  * Copyright (C) 2011 Google Inc.
      3  *
      4  * Licensed under the Apache License, Version 2.0 (the "License");
      5  * you may not use this file except in compliance with the License.
      6  * You may obtain a copy of the License at
      7  *
      8  *      http://www.apache.org/licenses/LICENSE-2.0
      9  *
     10  * Unless required by applicable law or agreed to in writing, software
     11  * distributed under the License is distributed on an "AS IS" BASIS,
     12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     13  * See the License for the specific language governing permissions and
     14  * limitations under the License.
     15  */
     16 package com.squareup.okhttp.mockwebserver;
     17 
     18 import com.squareup.okhttp.internal.Util;
     19 import java.io.ByteArrayInputStream;
     20 import java.io.ByteArrayOutputStream;
     21 import java.io.IOException;
     22 import java.io.InputStream;
     23 import java.io.UnsupportedEncodingException;
     24 import java.util.ArrayList;
     25 import java.util.Iterator;
     26 import java.util.List;
     27 import java.util.concurrent.TimeUnit;
     28 
     29 /** A scripted response to be replayed by the mock web server. */
     30 public final class MockResponse implements Cloneable {
     31   private static final String CHUNKED_BODY_HEADER = "Transfer-encoding: chunked";
     32 
     33   private String status = "HTTP/1.1 200 OK";
     34   private List<String> headers = new ArrayList<String>();
     35 
     36   /** The response body content, or null if {@code bodyStream} is set. */
     37   private byte[] body;
     38   /** The response body content, or null if {@code body} is set. */
     39   private InputStream bodyStream;
     40 
     41   private int throttleBytesPerPeriod = Integer.MAX_VALUE;
     42   private long throttlePeriod = 1;
     43   private TimeUnit throttleUnit = TimeUnit.SECONDS;
     44 
     45   private SocketPolicy socketPolicy = SocketPolicy.KEEP_OPEN;
     46 
     47   private int bodyDelayTimeMs = 0;
     48 
     49   private List<PushPromise> promises = new ArrayList<PushPromise>();
     50 
     51   /** Creates a new mock response with an empty body. */
     52   public MockResponse() {
     53     setBody(new byte[0]);
     54   }
     55 
     56   @Override public MockResponse clone() {
     57     try {
     58       MockResponse result = (MockResponse) super.clone();
     59       result.headers = new ArrayList<String>(headers);
     60       result.promises = new ArrayList<PushPromise>(promises);
     61       return result;
     62     } catch (CloneNotSupportedException e) {
     63       throw new AssertionError();
     64     }
     65   }
     66 
     67   /** Returns the HTTP response line, such as "HTTP/1.1 200 OK". */
     68   public String getStatus() {
     69     return status;
     70   }
     71 
     72   public MockResponse setResponseCode(int code) {
     73     this.status = "HTTP/1.1 " + code + " OK";
     74     return this;
     75   }
     76 
     77   public MockResponse setStatus(String status) {
     78     this.status = status;
     79     return this;
     80   }
     81 
     82   /** Returns the HTTP headers, such as "Content-Length: 0". */
     83   public List<String> getHeaders() {
     84     return headers;
     85   }
     86 
     87   /**
     88    * Removes all HTTP headers including any "Content-Length" and
     89    * "Transfer-encoding" headers that were added by default.
     90    */
     91   public MockResponse clearHeaders() {
     92     headers.clear();
     93     return this;
     94   }
     95 
     96   /**
     97    * Adds {@code header} as an HTTP header. For well-formed HTTP {@code header}
     98    * should contain a name followed by a colon and a value.
     99    */
    100   public MockResponse addHeader(String header) {
    101     headers.add(header);
    102     return this;
    103   }
    104 
    105   /**
    106    * Adds a new header with the name and value. This may be used to add multiple
    107    * headers with the same name.
    108    */
    109   public MockResponse addHeader(String name, Object value) {
    110     return addHeader(name + ": " + String.valueOf(value));
    111   }
    112 
    113   /**
    114    * Removes all headers named {@code name}, then adds a new header with the
    115    * name and value.
    116    */
    117   public MockResponse setHeader(String name, Object value) {
    118     removeHeader(name);
    119     return addHeader(name, value);
    120   }
    121 
    122   /** Removes all headers named {@code name}. */
    123   public MockResponse removeHeader(String name) {
    124     name += ":";
    125     for (Iterator<String> i = headers.iterator(); i.hasNext(); ) {
    126       String header = i.next();
    127       if (name.regionMatches(true, 0, header, 0, name.length())) {
    128         i.remove();
    129       }
    130     }
    131     return this;
    132   }
    133 
    134   /** Returns the raw HTTP payload, or null if this response is streamed. */
    135   public byte[] getBody() {
    136     return body;
    137   }
    138 
    139   /** Returns an input stream containing the raw HTTP payload. */
    140   InputStream getBodyStream() {
    141     return bodyStream != null ? bodyStream : new ByteArrayInputStream(body);
    142   }
    143 
    144   public MockResponse setBody(byte[] body) {
    145     setHeader("Content-Length", body.length);
    146     this.body = body;
    147     this.bodyStream = null;
    148     return this;
    149   }
    150 
    151   public MockResponse setBody(InputStream bodyStream, long bodyLength) {
    152     setHeader("Content-Length", bodyLength);
    153     this.body = null;
    154     this.bodyStream = bodyStream;
    155     return this;
    156   }
    157 
    158   /** Sets the response body to the UTF-8 encoded bytes of {@code body}. */
    159   public MockResponse setBody(String body) {
    160     try {
    161       return setBody(body.getBytes("UTF-8"));
    162     } catch (UnsupportedEncodingException e) {
    163       throw new AssertionError();
    164     }
    165   }
    166 
    167   /**
    168    * Sets the response body to {@code body}, chunked every {@code maxChunkSize}
    169    * bytes.
    170    */
    171   public MockResponse setChunkedBody(byte[] body, int maxChunkSize) {
    172     removeHeader("Content-Length");
    173     headers.add(CHUNKED_BODY_HEADER);
    174 
    175     try {
    176       ByteArrayOutputStream bytesOut = new ByteArrayOutputStream();
    177       int pos = 0;
    178       while (pos < body.length) {
    179         int chunkSize = Math.min(body.length - pos, maxChunkSize);
    180         bytesOut.write(Integer.toHexString(chunkSize).getBytes(Util.US_ASCII));
    181         bytesOut.write("\r\n".getBytes(Util.US_ASCII));
    182         bytesOut.write(body, pos, chunkSize);
    183         bytesOut.write("\r\n".getBytes(Util.US_ASCII));
    184         pos += chunkSize;
    185       }
    186       bytesOut.write("0\r\n\r\n".getBytes(Util.US_ASCII)); // Last chunk + empty trailer + crlf.
    187 
    188       this.body = bytesOut.toByteArray();
    189       return this;
    190     } catch (IOException e) {
    191       throw new AssertionError(); // In-memory I/O doesn't throw IOExceptions.
    192     }
    193   }
    194 
    195   /**
    196    * Sets the response body to the UTF-8 encoded bytes of {@code body}, chunked
    197    * every {@code maxChunkSize} bytes.
    198    */
    199   public MockResponse setChunkedBody(String body, int maxChunkSize) {
    200     try {
    201       return setChunkedBody(body.getBytes("UTF-8"), maxChunkSize);
    202     } catch (UnsupportedEncodingException e) {
    203       throw new AssertionError();
    204     }
    205   }
    206 
    207   public SocketPolicy getSocketPolicy() {
    208     return socketPolicy;
    209   }
    210 
    211   public MockResponse setSocketPolicy(SocketPolicy socketPolicy) {
    212     this.socketPolicy = socketPolicy;
    213     return this;
    214   }
    215 
    216   /**
    217    * Throttles the response body writer to sleep for the given period after each
    218    * series of {@code bytesPerPeriod} bytes are written. Use this to simulate
    219    * network behavior.
    220    */
    221   public MockResponse throttleBody(int bytesPerPeriod, long period, TimeUnit unit) {
    222     this.throttleBytesPerPeriod = bytesPerPeriod;
    223     this.throttlePeriod = period;
    224     this.throttleUnit = unit;
    225     return this;
    226   }
    227 
    228   public int getThrottleBytesPerPeriod() {
    229     return throttleBytesPerPeriod;
    230   }
    231 
    232   public long getThrottlePeriod() {
    233     return throttlePeriod;
    234   }
    235 
    236   public TimeUnit getThrottleUnit() {
    237     return throttleUnit;
    238   }
    239 
    240   /**
    241    * Set the delayed time of the response body to {@code delay}. This applies to the
    242    * response body only; response headers are not affected.
    243    */
    244   public MockResponse setBodyDelayTimeMs(int delay) {
    245     bodyDelayTimeMs = delay;
    246     return this;
    247   }
    248 
    249   public int getBodyDelayTimeMs() {
    250     return bodyDelayTimeMs;
    251   }
    252 
    253   /**
    254    * When {@link MockWebServer#setNpnProtocols(java.util.List) protocols}
    255    * include a SPDY variant, this attaches a pushed stream to this response.
    256    */
    257   public MockResponse withPush(PushPromise promise) {
    258     this.promises.add(promise);
    259     return this;
    260   }
    261 
    262   /** Returns the streams the server will push with this response. */
    263   public List<PushPromise> getPushPromises() {
    264     return promises;
    265   }
    266 
    267   @Override public String toString() {
    268     return status;
    269   }
    270 }
    271