Same session after mutate for WebTestClient mock server

Issue: SPR-15674

spring-test/src/main/java/org/springframework/test/web/reactive/server/AbstractMockServerSpec.java

@@ -20,8 +20,11 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 
+import org.springframework.util.Assert;
 import org.springframework.web.server.WebFilter;
 import org.springframework.web.server.adapter.WebHttpHandlerBuilder;
+import org.springframework.web.server.session.DefaultWebSessionManager;
+import org.springframework.web.server.session.WebSessionManager;
 
 /**
  * Base class for implementations of {@link WebTestClient.MockServerSpec}.
@@ -34,6 +37,8 @@ abstract class AbstractMockServerSpec<B extends WebTestClient.MockServerSpec<B>>
 
 	private final List<WebFilter> filters = new ArrayList<>(4);
 
+	private WebSessionManager sessionManager = new DefaultWebSessionManager();
+
 	private final List<MockServerConfigurer> configurers = new ArrayList<>(4);
 
 
@@ -43,6 +48,13 @@ abstract class AbstractMockServerSpec<B extends WebTestClient.MockServerSpec<B>>
 		return self();
 	}
 
+	@Override
+	public <T extends B> T webSessionManager(WebSessionManager sessionManager) {
+		Assert.notNull(sessionManager, "WebSessionManager must not be null.");
+		this.sessionManager = sessionManager;
+		return self();
+	}
+
 	@Override
 	public <T extends B> T apply(MockServerConfigurer configurer) {
 		configurer.afterConfigureAdded(this);
@@ -59,6 +71,7 @@ abstract class AbstractMockServerSpec<B extends WebTestClient.MockServerSpec<B>>
 	public WebTestClient.Builder configureClient() {
 		WebHttpHandlerBuilder builder = initHttpHandlerBuilder();
 		builder.filters(theFilters -> theFilters.addAll(0, this.filters));
+		builder.sessionManager(this.sessionManager);
 		this.configurers.forEach(configurer -> configurer.beforeServerCreated(builder));
 		return new DefaultWebTestClientBuilder(builder);
 	}

spring-test/src/main/java/org/springframework/test/web/reactive/server/WebTestClient.java

@@ -50,6 +50,7 @@ import org.springframework.web.reactive.function.server.RouterFunction;
 import org.springframework.web.reactive.result.method.annotation.ArgumentResolverConfigurer;
 import org.springframework.web.server.WebFilter;
 import org.springframework.web.server.WebHandler;
+import org.springframework.web.server.session.WebSessionManager;
 import org.springframework.web.util.UriBuilder;
 import org.springframework.web.util.UriBuilderFactory;
 
@@ -191,12 +192,20 @@ public interface WebTestClient {
 	interface MockServerSpec<B extends MockServerSpec<B>> {
 
 		/**
-		 * Register one or more {@link WebFilter} instances to apply to the
-		 * mock server.
+		 * Register {@link WebFilter} instances to add to the mock server.
 		 * @param filter one or more filters
 		 */
 		<T extends B> T webFilter(WebFilter... filter);
 
+		/**
+		 * Provide a session manager instance for the mock server.
+		 * <p>By default an instance of
+		 * {@link org.springframework.web.server.session.DefaultWebSessionManager
+		 * DefaultWebSessionManager} is used.
+		 * @param sessionManager the session manager to use
+		 */
+		<T extends B> T webSessionManager(WebSessionManager sessionManager);
+
 		/**
 		 * Shortcut for pre-packaged customizations to the mock server setup.
 		 * @param configurer the configurer to apply

spring-test/src/test/java/org/springframework/test/web/reactive/server/MockServerIntegrationTests.java

@@ -0,0 +1,71 @@
+/*
+ * 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 org.junit.Test;
+import reactor.core.publisher.Mono;
+
+import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DefaultDataBufferFactory;
+import org.springframework.http.ResponseCookie;
+import org.springframework.web.server.WebHandler;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+/**
+ * Mock server integration test scenarios.
+ * @author Rossen Stoyanchev
+ */
+public class MockServerIntegrationTests {
+
+
+	@Test
+	public void sameSessionInstanceAfterMutate() throws Exception {
+
+		WebHandler webHandler = exchange -> {
+			if (exchange.getRequest().getURI().getPath().equals("/set")) {
+				return exchange.getSession()
+						.doOnNext(session -> session.getAttributes().put("foo", "bar"))
+						.then();
+			}
+			else {
+				return exchange.getSession()
+						.map(session -> session.getAttribute("foo").orElse("none"))
+						.flatMap(value -> {
+							byte[] bytes = value.toString().getBytes(UTF_8);
+							DataBuffer buffer = new DefaultDataBufferFactory().wrap(bytes);
+							return exchange.getResponse().writeWith(Mono.just(buffer));
+						});
+			}
+		};
+
+		WebTestClient client = new DefaultMockServerSpec(webHandler).configureClient().build();
+
+		// Set the session attribute
+		EntityExchangeResult<Void> result = client.get().uri("/set").exchange()
+				.expectStatus().isOk().expectBody().isEmpty();
+
+		ResponseCookie session = result.getResponseCookies().getFirst("SESSION");
+
+		// Now get attribute
+		client.mutate().build()
+				.get().uri("/get")
+				.cookie(session.getName(), session.getValue())
+				.exchange()
+				.expectBody(String.class).isEqualTo("bar");
+	}
+
+}