1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17 package org.springframework.mock.web.test;
18
19 import java.io.ByteArrayOutputStream;
20 import java.io.IOException;
21 import java.io.OutputStream;
22 import java.io.OutputStreamWriter;
23 import java.io.PrintWriter;
24 import java.io.UnsupportedEncodingException;
25 import java.io.Writer;
26 import java.util.ArrayList;
27 import java.util.Collection;
28 import java.util.Collections;
29 import java.util.List;
30 import java.util.Locale;
31 import java.util.Map;
32 import javax.servlet.ServletOutputStream;
33 import javax.servlet.http.Cookie;
34 import javax.servlet.http.HttpServletResponse;
35
36 import org.springframework.util.Assert;
37 import org.springframework.util.LinkedCaseInsensitiveMap;
38 import org.springframework.web.util.WebUtils;
39
40
41
42
43
44
45
46
47
48
49
50
51 public class MockHttpServletResponse implements HttpServletResponse {
52
53 private static final String CHARSET_PREFIX = "charset=";
54
55 private static final String CONTENT_TYPE_HEADER = "Content-Type";
56
57 private static final String CONTENT_LENGTH_HEADER = "Content-Length";
58
59 private static final String LOCATION_HEADER = "Location";
60
61
62
63
64
65
66 private boolean outputStreamAccessAllowed = true;
67
68 private boolean writerAccessAllowed = true;
69
70 private String characterEncoding = WebUtils.DEFAULT_CHARACTER_ENCODING;
71
72 private boolean charset = false;
73
74 private final ByteArrayOutputStream content = new ByteArrayOutputStream();
75
76 private final ServletOutputStream outputStream = new ResponseServletOutputStream(this.content);
77
78 private PrintWriter writer;
79
80 private long contentLength = 0;
81
82 private String contentType;
83
84 private int bufferSize = 4096;
85
86 private boolean committed;
87
88 private Locale locale = Locale.getDefault();
89
90
91
92
93
94
95 private final List<Cookie> cookies = new ArrayList<Cookie>();
96
97 private final Map<String, HeaderValueHolder> headers = new LinkedCaseInsensitiveMap<HeaderValueHolder>();
98
99 private int status = HttpServletResponse.SC_OK;
100
101 private String errorMessage;
102
103 private String forwardedUrl;
104
105 private final List<String> includedUrls = new ArrayList<String>();
106
107
108
109
110
111
112
113
114
115
116 public void setOutputStreamAccessAllowed(boolean outputStreamAccessAllowed) {
117 this.outputStreamAccessAllowed = outputStreamAccessAllowed;
118 }
119
120
121
122
123 public boolean isOutputStreamAccessAllowed() {
124 return this.outputStreamAccessAllowed;
125 }
126
127
128
129
130
131 public void setWriterAccessAllowed(boolean writerAccessAllowed) {
132 this.writerAccessAllowed = writerAccessAllowed;
133 }
134
135
136
137
138 public boolean isWriterAccessAllowed() {
139 return this.writerAccessAllowed;
140 }
141
142 @Override
143 public void setCharacterEncoding(String characterEncoding) {
144 this.characterEncoding = characterEncoding;
145 this.charset = true;
146 updateContentTypeHeader();
147 }
148
149 private void updateContentTypeHeader() {
150 if (this.contentType != null) {
151 StringBuilder sb = new StringBuilder(this.contentType);
152 if (!this.contentType.toLowerCase().contains(CHARSET_PREFIX) && this.charset) {
153 sb.append(";").append(CHARSET_PREFIX).append(this.characterEncoding);
154 }
155 doAddHeaderValue(CONTENT_TYPE_HEADER, sb.toString(), true);
156 }
157 }
158
159 @Override
160 public String getCharacterEncoding() {
161 return this.characterEncoding;
162 }
163
164 @Override
165 public ServletOutputStream getOutputStream() {
166 if (!this.outputStreamAccessAllowed) {
167 throw new IllegalStateException("OutputStream access not allowed");
168 }
169 return this.outputStream;
170 }
171
172 @Override
173 public PrintWriter getWriter() throws UnsupportedEncodingException {
174 if (!this.writerAccessAllowed) {
175 throw new IllegalStateException("Writer access not allowed");
176 }
177 if (this.writer == null) {
178 Writer targetWriter = (this.characterEncoding != null ?
179 new OutputStreamWriter(this.content, this.characterEncoding) : new OutputStreamWriter(this.content));
180 this.writer = new ResponsePrintWriter(targetWriter);
181 }
182 return this.writer;
183 }
184
185 public byte[] getContentAsByteArray() {
186 flushBuffer();
187 return this.content.toByteArray();
188 }
189
190 public String getContentAsString() throws UnsupportedEncodingException {
191 flushBuffer();
192 return (this.characterEncoding != null) ?
193 this.content.toString(this.characterEncoding) : this.content.toString();
194 }
195
196 @Override
197 public void setContentLength(int contentLength) {
198 this.contentLength = contentLength;
199 doAddHeaderValue(CONTENT_LENGTH_HEADER, contentLength, true);
200 }
201
202 public int getContentLength() {
203 return (int) this.contentLength;
204 }
205
206 public void setContentLengthLong(long contentLength) {
207 this.contentLength = contentLength;
208 doAddHeaderValue(CONTENT_LENGTH_HEADER, contentLength, true);
209 }
210
211 public long getContentLengthLong() {
212 return this.contentLength;
213 }
214
215 @Override
216 public void setContentType(String contentType) {
217 this.contentType = contentType;
218 if (contentType != null) {
219 int charsetIndex = contentType.toLowerCase().indexOf(CHARSET_PREFIX);
220 if (charsetIndex != -1) {
221 String encoding = contentType.substring(charsetIndex + CHARSET_PREFIX.length());
222 this.characterEncoding = encoding;
223 this.charset = true;
224 }
225 updateContentTypeHeader();
226 }
227 }
228
229 @Override
230 public String getContentType() {
231 return this.contentType;
232 }
233
234 @Override
235 public void setBufferSize(int bufferSize) {
236 this.bufferSize = bufferSize;
237 }
238
239 @Override
240 public int getBufferSize() {
241 return this.bufferSize;
242 }
243
244 @Override
245 public void flushBuffer() {
246 setCommitted(true);
247 }
248
249 @Override
250 public void resetBuffer() {
251 if (isCommitted()) {
252 throw new IllegalStateException("Cannot reset buffer - response is already committed");
253 }
254 this.content.reset();
255 }
256
257 private void setCommittedIfBufferSizeExceeded() {
258 int bufSize = getBufferSize();
259 if (bufSize > 0 && this.content.size() > bufSize) {
260 setCommitted(true);
261 }
262 }
263
264 public void setCommitted(boolean committed) {
265 this.committed = committed;
266 }
267
268 @Override
269 public boolean isCommitted() {
270 return this.committed;
271 }
272
273 @Override
274 public void reset() {
275 resetBuffer();
276 this.characterEncoding = null;
277 this.contentLength = 0;
278 this.contentType = null;
279 this.locale = null;
280 this.cookies.clear();
281 this.headers.clear();
282 this.status = HttpServletResponse.SC_OK;
283 this.errorMessage = null;
284 }
285
286 @Override
287 public void setLocale(Locale locale) {
288 this.locale = locale;
289 }
290
291 @Override
292 public Locale getLocale() {
293 return this.locale;
294 }
295
296
297
298
299
300
301 @Override
302 public void addCookie(Cookie cookie) {
303 Assert.notNull(cookie, "Cookie must not be null");
304 this.cookies.add(cookie);
305 }
306
307 public Cookie[] getCookies() {
308 return this.cookies.toArray(new Cookie[this.cookies.size()]);
309 }
310
311 public Cookie getCookie(String name) {
312 Assert.notNull(name, "Cookie name must not be null");
313 for (Cookie cookie : this.cookies) {
314 if (name.equals(cookie.getName())) {
315 return cookie;
316 }
317 }
318 return null;
319 }
320
321 @Override
322 public boolean containsHeader(String name) {
323 return (HeaderValueHolder.getByName(this.headers, name) != null);
324 }
325
326
327
328
329
330
331 @Override
332 public Collection<String> getHeaderNames() {
333 return this.headers.keySet();
334 }
335
336
337
338
339
340
341
342
343
344
345 @Override
346 public String getHeader(String name) {
347 HeaderValueHolder header = HeaderValueHolder.getByName(this.headers, name);
348 return (header != null ? header.getStringValue() : null);
349 }
350
351
352
353
354
355
356
357
358
359 @Override
360 public List<String> getHeaders(String name) {
361 HeaderValueHolder header = HeaderValueHolder.getByName(this.headers, name);
362 if (header != null) {
363 return header.getStringValues();
364 }
365 else {
366 return Collections.emptyList();
367 }
368 }
369
370
371
372
373
374
375
376 public Object getHeaderValue(String name) {
377 HeaderValueHolder header = HeaderValueHolder.getByName(this.headers, name);
378 return (header != null ? header.getValue() : null);
379 }
380
381
382
383
384
385
386 public List<Object> getHeaderValues(String name) {
387 HeaderValueHolder header = HeaderValueHolder.getByName(this.headers, name);
388 if (header != null) {
389 return header.getValues();
390 }
391 else {
392 return Collections.emptyList();
393 }
394 }
395
396
397
398
399
400 @Override
401 public String encodeURL(String url) {
402 return url;
403 }
404
405
406
407
408
409
410
411
412
413 @Override
414 public String encodeRedirectURL(String url) {
415 return encodeURL(url);
416 }
417
418 @Override
419 public String encodeUrl(String url) {
420 return encodeURL(url);
421 }
422
423 @Override
424 public String encodeRedirectUrl(String url) {
425 return encodeRedirectURL(url);
426 }
427
428 @Override
429 public void sendError(int status, String errorMessage) throws IOException {
430 if (isCommitted()) {
431 throw new IllegalStateException("Cannot set error status - response is already committed");
432 }
433 this.status = status;
434 this.errorMessage = errorMessage;
435 setCommitted(true);
436 }
437
438 @Override
439 public void sendError(int status) throws IOException {
440 if (isCommitted()) {
441 throw new IllegalStateException("Cannot set error status - response is already committed");
442 }
443 this.status = status;
444 setCommitted(true);
445 }
446
447 @Override
448 public void sendRedirect(String url) throws IOException {
449 if (isCommitted()) {
450 throw new IllegalStateException("Cannot send redirect - response is already committed");
451 }
452 Assert.notNull(url, "Redirect URL must not be null");
453 setHeader(LOCATION_HEADER, url);
454 setStatus(HttpServletResponse.SC_MOVED_TEMPORARILY);
455 setCommitted(true);
456 }
457
458 public String getRedirectedUrl() {
459 return getHeader(LOCATION_HEADER);
460 }
461
462 @Override
463 public void setDateHeader(String name, long value) {
464 setHeaderValue(name, value);
465 }
466
467 @Override
468 public void addDateHeader(String name, long value) {
469 addHeaderValue(name, value);
470 }
471
472 @Override
473 public void setHeader(String name, String value) {
474 setHeaderValue(name, value);
475 }
476
477 @Override
478 public void addHeader(String name, String value) {
479 addHeaderValue(name, value);
480 }
481
482 @Override
483 public void setIntHeader(String name, int value) {
484 setHeaderValue(name, value);
485 }
486
487 @Override
488 public void addIntHeader(String name, int value) {
489 addHeaderValue(name, value);
490 }
491
492 private void setHeaderValue(String name, Object value) {
493 if (setSpecialHeader(name, value)) {
494 return;
495 }
496 doAddHeaderValue(name, value, true);
497 }
498
499 private void addHeaderValue(String name, Object value) {
500 if (setSpecialHeader(name, value)) {
501 return;
502 }
503 doAddHeaderValue(name, value, false);
504 }
505
506 private boolean setSpecialHeader(String name, Object value) {
507 if (CONTENT_TYPE_HEADER.equalsIgnoreCase(name)) {
508 setContentType((String) value);
509 return true;
510 }
511 else if (CONTENT_LENGTH_HEADER.equalsIgnoreCase(name)) {
512 setContentLength(Integer.parseInt((String) value));
513 return true;
514 }
515 else {
516 return false;
517 }
518 }
519
520 private void doAddHeaderValue(String name, Object value, boolean replace) {
521 HeaderValueHolder header = HeaderValueHolder.getByName(this.headers, name);
522 Assert.notNull(value, "Header value must not be null");
523 if (header == null) {
524 header = new HeaderValueHolder();
525 this.headers.put(name, header);
526 }
527 if (replace) {
528 header.setValue(value);
529 }
530 else {
531 header.addValue(value);
532 }
533 }
534
535 @Override
536 public void setStatus(int status) {
537 if(!this.isCommitted()) {
538 this.status = status;
539 }
540 }
541
542 @Override
543 public void setStatus(int status, String errorMessage) {
544 if(!this.isCommitted()) {
545 this.status = status;
546 this.errorMessage = errorMessage;
547 }
548 }
549
550 @Override
551 public int getStatus() {
552 return this.status;
553 }
554
555 public String getErrorMessage() {
556 return this.errorMessage;
557 }
558
559
560
561
562
563
564 public void setForwardedUrl(String forwardedUrl) {
565 this.forwardedUrl = forwardedUrl;
566 }
567
568 public String getForwardedUrl() {
569 return this.forwardedUrl;
570 }
571
572 public void setIncludedUrl(String includedUrl) {
573 this.includedUrls.clear();
574 if (includedUrl != null) {
575 this.includedUrls.add(includedUrl);
576 }
577 }
578
579 public String getIncludedUrl() {
580 int count = this.includedUrls.size();
581 if (count > 1) {
582 throw new IllegalStateException(
583 "More than 1 URL included - check getIncludedUrls instead: " + this.includedUrls);
584 }
585 return (count == 1 ? this.includedUrls.get(0) : null);
586 }
587
588 public void addIncludedUrl(String includedUrl) {
589 Assert.notNull(includedUrl, "Included URL must not be null");
590 this.includedUrls.add(includedUrl);
591 }
592
593 public List<String> getIncludedUrls() {
594 return this.includedUrls;
595 }
596
597
598
599
600
601
602 private class ResponseServletOutputStream extends DelegatingServletOutputStream {
603
604 public ResponseServletOutputStream(OutputStream out) {
605 super(out);
606 }
607
608 @Override
609 public void write(int b) throws IOException {
610 super.write(b);
611 super.flush();
612 setCommittedIfBufferSizeExceeded();
613 }
614
615 @Override
616 public void flush() throws IOException {
617 super.flush();
618 setCommitted(true);
619 }
620 }
621
622
623
624
625
626
627 private class ResponsePrintWriter extends PrintWriter {
628
629 public ResponsePrintWriter(Writer out) {
630 super(out, true);
631 }
632
633 @Override
634 public void write(char buf[], int off, int len) {
635 super.write(buf, off, len);
636 super.flush();
637 setCommittedIfBufferSizeExceeded();
638 }
639
640 @Override
641 public void write(String s, int off, int len) {
642 super.write(s, off, len);
643 super.flush();
644 setCommittedIfBufferSizeExceeded();
645 }
646
647 @Override
648 public void write(int c) {
649 super.write(c);
650 super.flush();
651 setCommittedIfBufferSizeExceeded();
652 }
653
654 @Override
655 public void flush() {
656 super.flush();
657 setCommitted(true);
658 }
659 }
660
661 }