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
17package com.google.mockwebserver;
18
19import static com.google.mockwebserver.MockWebServer.ASCII;
20import java.io.ByteArrayOutputStream;
21import java.io.IOException;
22import java.io.UnsupportedEncodingException;
23import java.util.ArrayList;
24import java.util.Iterator;
25import java.util.List;
26
27/**
28 * A scripted response to be replayed by the mock web server.
29 */
30public final class MockResponse implements Cloneable {
31    private static final String EMPTY_BODY_HEADER = "Content-Length: 0";
32    private static final String CHUNKED_BODY_HEADER = "Transfer-encoding: chunked";
33    private static final byte[] EMPTY_BODY = new byte[0];
34
35    private String status = "HTTP/1.1 200 OK";
36    private List<String> headers = new ArrayList<String>();
37    private byte[] body = EMPTY_BODY;
38    private int bytesPerSecond = Integer.MAX_VALUE;
39    private SocketPolicy socketPolicy = SocketPolicy.KEEP_OPEN;
40
41    public MockResponse() {
42        headers.add(EMPTY_BODY_HEADER);
43    }
44
45    @Override public MockResponse clone() {
46        try {
47            MockResponse result = (MockResponse) super.clone();
48            result.headers = new ArrayList<String>(result.headers);
49            return result;
50        } catch (CloneNotSupportedException e) {
51            throw new AssertionError();
52        }
53    }
54
55    /**
56     * Returns the HTTP response line, such as "HTTP/1.1 200 OK".
57     */
58    public String getStatus() {
59        return status;
60    }
61
62    public MockResponse setResponseCode(int code) {
63        this.status = "HTTP/1.1 " + code + " OK";
64        return this;
65    }
66
67    public MockResponse setStatus(String status) {
68        this.status = status;
69        return this;
70    }
71
72    /**
73     * Returns the HTTP headers, such as "Content-Length: 0".
74     */
75    public List<String> getHeaders() {
76        return headers;
77    }
78
79    public MockResponse clearHeaders() {
80        headers.clear();
81        return this;
82    }
83
84    public MockResponse addHeader(String header) {
85        headers.add(header);
86        return this;
87    }
88
89    public MockResponse addHeader(String name, Object value) {
90        return addHeader(name + ": " + String.valueOf(value));
91    }
92
93    public MockResponse setHeader(String name, Object value) {
94        removeHeader(name);
95        return addHeader(name, value);
96    }
97
98    public MockResponse removeHeader(String name) {
99        name += ": ";
100        for (Iterator<String> i = headers.iterator(); i.hasNext();) {
101            String header = i.next();
102            if (name.regionMatches(true, 0, header, 0, name.length())) {
103                i.remove();
104            }
105        }
106        return this;
107    }
108
109    /**
110     * Returns an input stream containing the raw HTTP payload.
111     */
112    public byte[] getBody() {
113        return body;
114    }
115
116    public MockResponse setBody(byte[] body) {
117        if (this.body == EMPTY_BODY) {
118            headers.remove(EMPTY_BODY_HEADER);
119        }
120        this.headers.add("Content-Length: " + body.length);
121        this.body = body;
122        return this;
123    }
124
125    public MockResponse setBody(String body) {
126        try {
127            return setBody(body.getBytes(ASCII));
128        } catch (UnsupportedEncodingException e) {
129            throw new AssertionError();
130        }
131    }
132
133    public MockResponse setChunkedBody(byte[] body, int maxChunkSize) throws IOException {
134        headers.remove(EMPTY_BODY_HEADER);
135        headers.add(CHUNKED_BODY_HEADER);
136
137        ByteArrayOutputStream bytesOut = new ByteArrayOutputStream();
138        int pos = 0;
139        while (pos < body.length) {
140            int chunkSize = Math.min(body.length - pos, maxChunkSize);
141            bytesOut.write(Integer.toHexString(chunkSize).getBytes(ASCII));
142            bytesOut.write("\r\n".getBytes(ASCII));
143            bytesOut.write(body, pos, chunkSize);
144            bytesOut.write("\r\n".getBytes(ASCII));
145            pos += chunkSize;
146        }
147        bytesOut.write("0\r\n\r\n".getBytes(ASCII)); // last chunk + empty trailer + crlf
148        this.body = bytesOut.toByteArray();
149        return this;
150    }
151
152    public MockResponse setChunkedBody(String body, int maxChunkSize) throws IOException {
153        return setChunkedBody(body.getBytes(ASCII), maxChunkSize);
154    }
155
156    public SocketPolicy getSocketPolicy() {
157        return socketPolicy;
158    }
159
160    public MockResponse setSocketPolicy(SocketPolicy socketPolicy) {
161        this.socketPolicy = socketPolicy;
162        return this;
163    }
164
165    public int getBytesPerSecond() {
166        return bytesPerSecond;
167    }
168
169    /**
170     * Set simulated network speed, in bytes per second.
171     */
172    public MockResponse setBytesPerSecond(int bytesPerSecond) {
173        this.bytesPerSecond = bytesPerSecond;
174        return this;
175    }
176
177    @Override public String toString() {
178        return status;
179    }
180}
181