001/*
002 * Stallion Core: A Modern Web Framework
003 *
004 * Copyright (C) 2015 - 2016 Stallion Software LLC.
005 *
006 * This program is free software: you can redistribute it and/or modify it under the terms of the
007 * GNU General Public License as published by the Free Software Foundation, either version 2 of
008 * the License, or (at your option) any later version. This program is distributed in the hope that
009 * it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
010 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
011 * License for more details. You should have received a copy of the GNU General Public License
012 * along with this program.  If not, see <http://www.gnu.org/licenses/gpl-2.0.html>.
013 *
014 *
015 *
016 */
017
018package io.stallion.testing;
019
020import io.stallion.requests.StResponse;
021import io.stallion.utils.json.JSON;
022
023import javax.servlet.ServletOutputStream;
024import javax.servlet.http.Cookie;
025import java.io.IOException;
026import java.io.PrintWriter;
027import java.io.StringWriter;
028import java.util.HashMap;
029import java.util.List;
030import java.util.Map;
031
032import static io.stallion.utils.Literals.*;
033
034public class MockResponse<T> extends StResponse {
035
036    private String contentType;
037    private int status = 0;
038    private int contentLength = 0;
039    private StringWriter stringWriter;
040    private PrintWriter printWriter;
041    private Map<String, String> headers;
042    private ServletOutputStream outputStream;
043    private List<Cookie> cookies = list();
044    private boolean hasByteOutput = false;
045
046    public MockResponse() {
047        stringWriter = new StringWriter();
048        printWriter = new PrintWriter(stringWriter);
049        headers = new HashMap<String, String>();
050        outputStream = new MockOutputStream();
051    }
052
053    @Override
054    public StResponse setDefaultContentType(String contentType) {
055        return this;
056    }
057
058    public StResponse setContentType(String contentType) {
059        this.contentType = contentType;
060        return this;
061    }
062
063    public StResponse setStatus(int sc) {
064        this.status = sc;
065        return this;
066    }
067
068    public int getStatus() {
069        return this.status;
070    }
071
072    public PrintWriter getWriter() throws IOException {
073        return printWriter;
074    }
075
076    public String getContent() {
077        String out = outputStream.toString();
078        if (!empty(out)) {
079            return out;
080        }
081        return stringWriter.toString();
082    }
083
084    public StResponse addHeader(String name, String value) {
085        headers.put(name, value);
086        return this;
087    }
088
089    public ServletOutputStream getOutputStream() throws IOException {
090        return outputStream;
091    }
092
093    public StResponse setDateHeader(String name, Long value) {
094        addHeader(name, value.toString());
095        return this;
096    }
097
098    public void setContentLength(int length) {
099        this.contentLength = length;
100    }
101
102    public <T> T asObject(Class<? extends T> cls) {
103        try {
104            return JSON.parse(this.getContent(), cls);
105        } catch (Exception e) {
106            throw new RuntimeException(e);
107        }
108    }
109
110    public Map asMap() {
111        return JSON.parseMap(getContent());
112    }
113
114    public String getHeader(String name) {
115        return headers.getOrDefault(name, null);
116    }
117
118    public Cookie addCookie(Cookie cookie) {
119        cookies.add(cookie);
120        return cookie;
121    }
122
123    public Cookie getCookie(String name) {
124        for (Cookie cookie: cookies) {
125            if (cookie.getName().equals(name)) {
126                return cookie;
127            }
128        }
129        return null;
130    }
131
132}