/*
 * Copyright 2002-2017 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.springframework.test.web.reactive.server;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import java.util.function.UnaryOperator;

import reactor.core.publisher.Mono;

import org.springframework.http.HttpHeaders;
import org.springframework.util.Assert;
import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;


/**
 * Apply {@code ServerWebExchange} transformations during "mock" server tests
 * with the {@code WebTestClient}.
 *
 * <p>Register the {@code WebFilter} while setting up the mock server through
 * one of the following:
 * <ul>
 * <li>{@link WebTestClient#bindToController}
 * <li>{@link WebTestClient#bindToRouterFunction}
 * <li>{@link WebTestClient#bindToApplicationContext}
 * <li>{@link WebTestClient#bindToWebHandler}
 * </ul>
 *
 * <p>Example usage:
 *
 * <pre class="code">
 * ExchangeMutatorWebFilter mutator = new ExchangeMutatorWebFilter(exchange -> ...);
 * WebTestClient client = WebTestClient.bindToController(new MyController()).webFilter(mutator).build();
 * </pre>
 *
 * <p>It is also possible to apply "per request" transformations:
 *
 * <pre class="code">
 * ExchangeMutatorWebFilter mutator = new ExchangeMutatorWebFilter(exchange -> ...);
 * WebTestClient client = WebTestClient.bindToController(new MyController()).webFilter(mutator).build();
 * client.filter(mutator.perClient(exchange -> ...)).get().uri("/").exchange();
 * </pre>
 *
 * @author Rossen Stoyanchev
 * @since 5.0
 */
public class ExchangeMutatorWebFilter implements WebFilter {

	private final Function<ServerWebExchange, ServerWebExchange> processor;

	private final Map<String, Function<ServerWebExchange, ServerWebExchange>> perRequestProcessors =
			new ConcurrentHashMap<>(4);


	public ExchangeMutatorWebFilter(UnaryOperator<ServerWebExchange> processor) {
		Assert.notNull(processor, "'processor' is required");
		this.processor = processor;
	}


	@Override
	public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
		exchange = getProcessor(exchange).apply(exchange);
		return chain.filter(exchange);
	}

	private Function<ServerWebExchange, ServerWebExchange> getProcessor(ServerWebExchange exchange) {
		String id = getRequestId(exchange.getRequest().getHeaders());
		Function<ServerWebExchange, ServerWebExchange> clientMutator = this.perRequestProcessors.remove(id);
		return (clientMutator != null ? this.processor.andThen(clientMutator) : this.processor);
	}

	private String getRequestId(HttpHeaders headers) {
		String id = headers.getFirst(WebTestClient.WEBTESTCLIENT_REQUEST_ID);
		Assert.notNull(id, "No \"" + WebTestClient.WEBTESTCLIENT_REQUEST_ID + "\" header");
		return id;
	}

	/**
	 * Apply the given processor only to requests performed through the client
	 * instance filtered with the returned filter. See class-level Javadoc for
	 * sample code.
	 * @param processor the exchange processor to use
	 * @return client filter for use with {@link WebTestClient#filter}
	 */
	public ExchangeFilterFunction perClient(UnaryOperator<ServerWebExchange> processor) {
		return (request, next) -> {
			String id = getRequestId(request.headers());
			this.perRequestProcessors.compute(id,
					(s, value) -> value != null ? value.andThen(processor) : processor);
			return next.exchange(request);
		};
	}

}

@@ -777,49 +777,8 @@
 ion;
-%0Aimport java.util.function.UnaryOperator;
 %0A%0Aim
@@ -1616,32 +1616,169 @@
 e class=%22code%22%3E%0A
+ * Function&lt;ServerWebExchange, ServerWebExchange&gt; fn1 = ...;%0A * Function&lt;ServerWebExchange, ServerWebExchange&gt; fn2 = ...;%0A *%0A
  * ExchangeMutat
@@ -1824,39 +1824,44 @@
 orWebFilter(
-exchange -%3E ...
+fn1().andThen(fn2())
 );%0A * WebTes
@@ -1965,24 +1965,27 @@
 * %3C/pre%3E%0A *%0A
+ *%0A
  * %3Cp%3EIt is 
@@ -2130,31 +2130,16 @@
 bFilter(
-exchange -%3E ...
 );%0A * We
@@ -2238,16 +2238,156 @@
 uild();%0A
+ *%0A * Function&lt;ServerWebExchange, ServerWebExchange&gt; fn1 = ...;%0A * Function&lt;ServerWebExchange, ServerWebExchange&gt; fn2 = ...;%0A *%0A
  * clien
@@ -2417,23 +2417,28 @@
 ent(
-exchange -%3E ...
+fn1().andThen(fn2())
 )).g
@@ -2823,30 +2823,115 @@
 bFilter(
-UnaryOperator%3C
+) %7B%0A%09%09this(exchange -%3E exchange);%0A%09%7D%0A%0A%09public ExchangeMutatorWebFilter(Function%3CServerWebExchange, 
 ServerWe
@@ -4155,22 +4155,36 @@
 ent(
-UnaryOperator%3C
+Function%3CServerWebExchange, 
 Serv

Fix composability issue in ExchangeMutatorWebFilter