001/*
002 * Copyright 2002-2020 the original author or authors.
003 *
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 *
008 *      https://www.apache.org/licenses/LICENSE-2.0
009 *
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 */
016
017package org.springframework.mock.http.client.reactive;
018
019import java.nio.ByteBuffer;
020import java.nio.charset.Charset;
021import java.nio.charset.StandardCharsets;
022import java.util.Collection;
023
024import org.reactivestreams.Publisher;
025import reactor.core.publisher.Flux;
026import reactor.core.publisher.Mono;
027
028import org.springframework.core.io.buffer.DataBuffer;
029import org.springframework.core.io.buffer.DataBufferFactory;
030import org.springframework.core.io.buffer.DataBufferUtils;
031import org.springframework.core.io.buffer.DefaultDataBufferFactory;
032import org.springframework.http.HttpHeaders;
033import org.springframework.http.HttpStatus;
034import org.springframework.http.MediaType;
035import org.springframework.http.ResponseCookie;
036import org.springframework.http.client.reactive.ClientHttpResponse;
037import org.springframework.util.Assert;
038import org.springframework.util.LinkedMultiValueMap;
039import org.springframework.util.MultiValueMap;
040
041/**
042 * Mock implementation of {@link ClientHttpResponse}.
043 *
044 * @author Brian Clozel
045 * @author Rossen Stoyanchev
046 * @since 5.0
047 */
048public class MockClientHttpResponse implements ClientHttpResponse {
049
050        private final int status;
051
052        private final HttpHeaders headers = new HttpHeaders();
053
054        private final MultiValueMap<String, ResponseCookie> cookies = new LinkedMultiValueMap<>();
055
056        private Flux<DataBuffer> body = Flux.empty();
057
058        private final DataBufferFactory bufferFactory = new DefaultDataBufferFactory();
059
060
061        public MockClientHttpResponse(HttpStatus status) {
062                Assert.notNull(status, "HttpStatus is required");
063                this.status = status.value();
064        }
065
066        public MockClientHttpResponse(int status) {
067                Assert.isTrue(status > 99 && status < 1000, "Status must be between 100 and 999");
068                this.status = status;
069        }
070
071
072        @Override
073        public HttpStatus getStatusCode() {
074                return HttpStatus.valueOf(this.status);
075        }
076
077        @Override
078        public int getRawStatusCode() {
079                return this.status;
080        }
081
082        @Override
083        public HttpHeaders getHeaders() {
084                if (!getCookies().isEmpty() && this.headers.get(HttpHeaders.SET_COOKIE) == null) {
085                        getCookies().values().stream().flatMap(Collection::stream)
086                                        .forEach(cookie -> getHeaders().add(HttpHeaders.SET_COOKIE, cookie.toString()));
087                }
088                return this.headers;
089        }
090
091        @Override
092        public MultiValueMap<String, ResponseCookie> getCookies() {
093                return this.cookies;
094        }
095
096        public void setBody(Publisher<DataBuffer> body) {
097                this.body = Flux.from(body);
098        }
099
100        public void setBody(String body) {
101                setBody(body, StandardCharsets.UTF_8);
102        }
103
104        public void setBody(String body, Charset charset) {
105                DataBuffer buffer = toDataBuffer(body, charset);
106                this.body = Flux.just(buffer);
107        }
108
109        private DataBuffer toDataBuffer(String body, Charset charset) {
110                byte[] bytes = body.getBytes(charset);
111                ByteBuffer byteBuffer = ByteBuffer.wrap(bytes);
112                return this.bufferFactory.wrap(byteBuffer);
113        }
114
115        @Override
116        public Flux<DataBuffer> getBody() {
117                return this.body;
118        }
119
120        /**
121         * Return the response body aggregated and converted to a String using the
122         * charset of the Content-Type response or otherwise as "UTF-8".
123         */
124        public Mono<String> getBodyAsString() {
125                return DataBufferUtils.join(getBody())
126                                .map(buffer -> {
127                                        String s = buffer.toString(getCharset());
128                                        DataBufferUtils.release(buffer);
129                                        return s;
130                                })
131                                .defaultIfEmpty("");
132        }
133
134        private Charset getCharset() {
135                Charset charset = null;
136                MediaType contentType = getHeaders().getContentType();
137                if (contentType != null) {
138                        charset = contentType.getCharset();
139                }
140                return (charset != null ? charset : StandardCharsets.UTF_8);
141        }
142
143
144        @Override
145        public String toString() {
146                HttpStatus code = HttpStatus.resolve(this.status);
147                return (code != null ? code.name() + "(" + this.status + ")" : "Status (" + this.status + ")") + this.headers;
148        }
149}