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.net.URI;
020import java.nio.charset.Charset;
021import java.nio.charset.StandardCharsets;
022import java.util.Collection;
023import java.util.Optional;
024import java.util.function.Function;
025
026import org.reactivestreams.Publisher;
027import reactor.core.publisher.Flux;
028import reactor.core.publisher.Mono;
029
030import org.springframework.core.io.buffer.DataBuffer;
031import org.springframework.core.io.buffer.DataBufferFactory;
032import org.springframework.core.io.buffer.DataBufferUtils;
033import org.springframework.core.io.buffer.DefaultDataBufferFactory;
034import org.springframework.http.HttpHeaders;
035import org.springframework.http.HttpMethod;
036import org.springframework.http.client.reactive.AbstractClientHttpRequest;
037import org.springframework.http.client.reactive.ClientHttpRequest;
038import org.springframework.util.Assert;
039import org.springframework.util.MimeType;
040import org.springframework.web.util.UriComponentsBuilder;
041
042/**
043 * Mock implementation of {@link ClientHttpRequest}.
044 *
045 * @author Brian Clozel
046 * @author Rossen Stoyanchev
047 * @since 5.0
048 */
049public class MockClientHttpRequest extends AbstractClientHttpRequest {
050
051        private final HttpMethod httpMethod;
052
053        private final URI url;
054
055        private final DataBufferFactory bufferFactory = new DefaultDataBufferFactory();
056
057        private Flux<DataBuffer> body = Flux.error(
058                        new IllegalStateException("The body is not set. " +
059                                        "Did handling complete with success? Is a custom \"writeHandler\" configured?"));
060
061        private Function<Flux<DataBuffer>, Mono<Void>> writeHandler;
062
063
064        public MockClientHttpRequest(HttpMethod httpMethod, String urlTemplate, Object... vars) {
065                this(httpMethod, UriComponentsBuilder.fromUriString(urlTemplate).buildAndExpand(vars).encode().toUri());
066        }
067
068        public MockClientHttpRequest(HttpMethod httpMethod, URI url) {
069                this.httpMethod = httpMethod;
070                this.url = url;
071                this.writeHandler = body -> {
072                        this.body = body.cache();
073                        return this.body.then();
074                };
075        }
076
077
078        /**
079         * Configure a custom handler for writing the request body.
080         *
081         * <p>The default write handler consumes and caches the request body so it
082         * may be accessed subsequently, e.g. in test assertions. Use this property
083         * when the request body is an infinite stream.
084         *
085         * @param writeHandler the write handler to use returning {@code Mono<Void>}
086         * when the body has been "written" (i.e. consumed).
087         */
088        public void setWriteHandler(Function<Flux<DataBuffer>, Mono<Void>> writeHandler) {
089                Assert.notNull(writeHandler, "'writeHandler' is required");
090                this.writeHandler = writeHandler;
091        }
092
093
094        @Override
095        public HttpMethod getMethod() {
096                return this.httpMethod;
097        }
098
099        @Override
100        public URI getURI() {
101                return this.url;
102        }
103
104        @Override
105        public DataBufferFactory bufferFactory() {
106                return this.bufferFactory;
107        }
108
109        @Override
110        protected void applyHeaders() {
111        }
112
113        @Override
114        protected void applyCookies() {
115                getCookies().values().stream().flatMap(Collection::stream)
116                                .forEach(cookie -> getHeaders().add(HttpHeaders.COOKIE, cookie.toString()));
117        }
118
119        @Override
120        public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
121                return doCommit(() -> Mono.defer(() -> this.writeHandler.apply(Flux.from(body))));
122        }
123
124        @Override
125        public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
126                return writeWith(Flux.from(body).flatMap(p -> p));
127        }
128
129        @Override
130        public Mono<Void> setComplete() {
131                return writeWith(Flux.empty());
132        }
133
134
135        /**
136         * Return the request body, or an error stream if the body was never set
137         * or when {@link #setWriteHandler} is configured.
138         */
139        public Flux<DataBuffer> getBody() {
140                return this.body;
141        }
142
143        /**
144         * Aggregate response data and convert to a String using the "Content-Type"
145         * charset or "UTF-8" by default.
146         */
147        public Mono<String> getBodyAsString() {
148
149                Charset charset = Optional.ofNullable(getHeaders().getContentType()).map(MimeType::getCharset)
150                                .orElse(StandardCharsets.UTF_8);
151
152                return DataBufferUtils.join(getBody())
153                                .map(buffer -> {
154                                        String s = buffer.toString(charset);
155                                        DataBufferUtils.release(buffer);
156                                        return s;
157                                })
158                                .defaultIfEmpty("");
159        }
160
161}