001/**
002 * Copyright (c) 2015-2022, Michael Yang 杨福海 (fuhai999@gmail.com).
003 * <p>
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 * <p>
008 * http://www.apache.org/licenses/LICENSE-2.0
009 * <p>
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016package io.jboot.test;
017
018import io.jboot.components.http.HttpMimeTypes;
019import io.jboot.test.web.MockHttpServletRequest;
020import io.jboot.test.web.MockHttpServletResponse;
021import io.jboot.test.web.MockServletInputStream;
022import io.jboot.utils.StrUtil;
023
024import javax.servlet.ServletInputStream;
025import javax.servlet.http.Cookie;
026import java.io.*;
027import java.nio.charset.StandardCharsets;
028import java.util.*;
029import java.util.function.Consumer;
030
031public class MockMvc {
032
033    protected boolean holdCookiesEnable = false;
034    protected Set<Cookie> holdCookies = new HashSet<>();
035
036    protected Consumer<MockHttpServletRequest> requestStartListener;
037    protected Consumer<MockHttpServletResponse> requestFinishedListener;
038
039    public MockMvc() {
040    }
041
042    public MockMvc(boolean holdCookiesEnable) {
043        this.holdCookiesEnable = holdCookiesEnable;
044    }
045
046    public boolean isHoldCookiesEnable() {
047        return holdCookiesEnable;
048    }
049
050    public void setHoldCookiesEnable(boolean holdCookiesEnable) {
051        this.holdCookiesEnable = holdCookiesEnable;
052    }
053
054    public Set<Cookie> getHoldCookies() {
055        return holdCookies;
056    }
057
058    public String getCookieValue(String name) {
059        for (Cookie holdCookie : holdCookies) {
060            if (holdCookie.getName().equals(name)) {
061                return holdCookie.getValue();
062            }
063        }
064        return null;
065    }
066
067    public void setHoldCookies(Set<Cookie> holdCookies) {
068        this.holdCookies = holdCookies;
069    }
070
071    public Consumer<MockHttpServletRequest> getRequestStartListener() {
072        return requestStartListener;
073    }
074
075    public void setRequestStartListener(Consumer<MockHttpServletRequest> requestStartListener) {
076        this.requestStartListener = requestStartListener;
077    }
078
079    public Consumer<MockHttpServletResponse> getRequestFinishedListener() {
080        return requestFinishedListener;
081    }
082
083    public void setRequestFinishedListener(Consumer<MockHttpServletResponse> requestFinishedListener) {
084        this.requestFinishedListener = requestFinishedListener;
085    }
086
087
088    public MockMvcResult get(String target) {
089        return get(target, null, null, null);
090    }
091
092
093    public MockMvcResult get(String target, Map<String, Object> paras) {
094        return get(target, paras, null, null);
095    }
096
097
098    public MockMvcResult get(String target, Map<String, Object> p, Set<Cookie> cookies) {
099        return get(target, p, null, cookies);
100    }
101
102
103    public MockMvcResult get(String target, Map<String, Object> p, Map<String, String> headers) {
104        return get(target, p, headers, null);
105    }
106
107    public MockMvcResult get(String target, Map<String, Object> p, Map<String, String> headers, Set<Cookie> cookies) {
108        MockHttpServletRequest request = new MockHttpServletRequest();
109        request.setMethod("GET");
110
111        final Map<String, Object> paras = p != null ? p : new HashMap<>();
112
113
114        int indexOf = target.lastIndexOf("?");
115        if (indexOf != -1) {
116            Map<String, String> targetParas = StrUtil.queryStringToMap(target.substring(indexOf + 1));
117            paras.putAll(targetParas);
118            target = target.substring(0, indexOf);
119        }
120
121        request.setServletPath(target);
122
123        if (headers != null) {
124            request.setHeaders(headers);
125        }
126
127        paras.forEach(request::addParameter);
128
129        if (cookies != null) {
130            request.setCookies(cookies);
131        }
132
133
134        return doMockRequest(request);
135    }
136
137    public MockMvcResult post(String target) {
138        return post(target, null, null, null, null);
139    }
140
141    public MockMvcResult post(String target, String postData) {
142        return post(target, null, null, null, postData);
143    }
144
145    public MockMvcResult post(String target, Map<String, Object> paras) {
146        return post(target, paras, null, null, null);
147    }
148
149    public MockMvcResult post(String target, Map<String, Object> paras, String postData) {
150        return post(target, paras, null, null, postData);
151    }
152
153    public MockMvcResult post(String target, Map<String, Object> paras, Map<String, String> headers) {
154        return post(target, paras, headers, null, null);
155    }
156
157    public MockMvcResult post(String target, Map<String, Object> paras, Map<String, String> headers, String postData) {
158        return post(target, paras, headers, null, postData);
159    }
160
161    public MockMvcResult post(String target, Map<String, Object> paras, Map<String, String> headers, Set<Cookie> cookies, String postData) {
162        MockServletInputStream inStream = null;
163        if (StrUtil.isNotBlank(postData)) {
164            inStream = new MockServletInputStream(postData);
165        }
166        return doPost(target, paras, headers, cookies, inStream);
167    }
168
169
170    public MockMvcResult upload(String target, Map<String, Object> paras, Map<String, String> headers, Set<Cookie> cookies) {
171
172
173        String endFlag = "\r\n";
174        String startFlag = "--";
175        String boundary = "------" + StrUtil.uuid();
176
177        if (headers == null) {
178            headers = new HashMap<>();
179        }
180        headers.put("Content-Type", "multipart/form-data; boundary=" + boundary);
181
182        ByteArrayOutputStream baos = new ByteArrayOutputStream();
183        if (paras != null) {
184            paras.forEach((key, value) -> {
185                if (value instanceof File) {
186                    File file = (File) value;
187                    writeString(baos, startFlag + boundary + endFlag);
188                    writeString(baos, "Content-Disposition: form-data; name=\"" + key + "\"; filename=\"" + file.getName() + "\"");
189                    writeString(baos, endFlag + "Content-Type: " + HttpMimeTypes.getMimeType(file.getName()));
190                    writeString(baos, endFlag + endFlag);
191                    writeFile(baos, file);
192                    writeString(baos, endFlag);
193                } else {
194                    writeString(baos, startFlag + boundary + endFlag);
195                    writeString(baos, "Content-Disposition: form-data; name=\"" + key + "\"");
196                    writeString(baos, endFlag + endFlag);
197                    writeString(baos, String.valueOf(value));
198                    writeString(baos, endFlag);
199                }
200            });
201        }
202
203        writeString(baos, startFlag + boundary + startFlag + endFlag);
204        return doPost(target, paras, headers, cookies, new MockServletInputStream(baos.toByteArray()));
205    }
206
207    private void writeFile(ByteArrayOutputStream dos, File file) {
208        FileInputStream fStream = null;
209        try {
210            fStream = new FileInputStream(file);
211            byte[] buffer = new byte[2028];
212            for (int len = 0; (len = fStream.read(buffer)) > 0; ) {
213                dos.write(buffer, 0, len);
214            }
215        } catch (IOException e) {
216            e.printStackTrace();
217        } finally {
218            if (fStream != null) {
219                try {
220                    fStream.close();
221                } catch (IOException e) {
222                }
223            }
224        }
225
226    }
227
228    private void writeString(OutputStream dos, String s) {
229        try {
230            dos.write(s.getBytes(StandardCharsets.UTF_8));
231        } catch (IOException e) {
232            e.printStackTrace();
233        }
234    }
235
236
237    public MockMvcResult doPost(String target, Map<String, Object> paras, Map<String, String> headers, Set<Cookie> cookies, ServletInputStream inStream) {
238        MockHttpServletRequest request = new MockHttpServletRequest();
239        request.setMethod("POST");
240
241        int indexOf = target.lastIndexOf("?");
242        if (indexOf != -1) {
243            Map<String, String> targetParas = StrUtil.queryStringToMap(target.substring(indexOf + 1));
244            if (paras == null) {
245                paras = new HashMap<>();
246            }
247            paras.putAll(targetParas);
248            target = target.substring(0, indexOf);
249        }
250
251        request.setServletPath(target);
252
253        if (headers != null) {
254            request.setHeaders(headers);
255        }
256
257        if (paras != null) {
258            paras.forEach(request::addParameter);
259        }
260
261        if (inStream != null) {
262            request.setInputStream(inStream);
263        }
264
265
266        if (cookies != null) {
267            request.setCookies(cookies);
268        }
269
270        return doMockRequest(request);
271    }
272
273
274    public MockMvcResult doMockRequest(MockHttpServletRequest request) {
275        MockHttpServletResponse response = createResponse();
276        try {
277            if (requestStartListener != null) {
278                requestStartListener.accept(request);
279            }
280
281            Set<Cookie> cookies = new HashSet<>();
282
283            //开启 cookie 保持, 用户未设置自己的 cookie, 上次的 cookie 有值
284            if (isHoldCookiesEnable()) {
285                cookies.addAll(holdCookies);
286            }
287
288            for (Cookie cookie : request.getCookies()) {
289                doSetCookie(cookie, cookies);
290            }
291
292            request.setCookies(cookies);
293            doSendRequest(request, response);
294
295        } finally {
296            if (requestFinishedListener != null) {
297                requestFinishedListener.accept(response);
298            }
299
300            if (isHoldCookiesEnable() && response.getCookies().size() > 0) {
301                response.getCookies().forEach(cookie -> doSetCookie(cookie, holdCookies));
302            }
303        }
304
305        return new MockMvcResult(response);
306    }
307
308    public void doSetCookie(Cookie newCookie, Set<Cookie> toCookies) {
309        toCookies.removeIf(cookie -> Objects.equals(cookie.getName(), newCookie.getName()));
310        if (StrUtil.isNotEmpty(newCookie.getValue())) {
311            toCookies.add(newCookie);
312        }
313    }
314
315
316    public void doSendRequest(MockHttpServletRequest request, MockHttpServletResponse response) {
317        MockApp.mockRequest(request, response);
318    }
319
320    public MockHttpServletResponse createResponse() {
321        return new MockHttpServletResponse();
322    }
323
324}