MockClientHttpRequest|Response support cookie headers

Issue: SPR-16124

spring-test/src/main/java/org/springframework/mock/http/client/reactive/MockClientHttpRequest.java

@@ -17,6 +17,7 @@
 package org.springframework.mock.http.client.reactive;
 
 import java.net.URI;
+import java.util.Collection;
 import java.util.function.Function;
 
 import org.reactivestreams.Publisher;
@@ -26,6 +27,7 @@ import reactor.core.publisher.Mono;
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
 import org.springframework.core.io.buffer.DefaultDataBufferFactory;
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.client.reactive.AbstractClientHttpRequest;
 import org.springframework.http.client.reactive.ClientHttpRequest;
@@ -114,6 +116,8 @@ public class MockClientHttpRequest extends AbstractClientHttpRequest {
 
 	@Override
 	protected void applyCookies() {
+		getCookies().values().stream().flatMap(Collection::stream)
+				.forEach(cookie -> getHeaders().add(HttpHeaders.COOKIE, cookie.toString()));
 	}
 
 	@Override

spring-test/src/main/java/org/springframework/mock/http/client/reactive/MockClientHttpResponse.java

@@ -19,6 +19,7 @@ package org.springframework.mock.http.client.reactive;
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
+import java.util.Collection;
 
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
@@ -62,15 +63,22 @@ public class MockClientHttpResponse implements ClientHttpResponse {
 	}
 
 
+	@Override
 	public HttpStatus getStatusCode() {
 		return this.status;
 	}
 
 	@Override
 	public HttpHeaders getHeaders() {
+		String headerName = HttpHeaders.SET_COOKIE;
+		if (!getCookies().isEmpty() && this.headers.get(headerName) == null) {
+			getCookies().values().stream().flatMap(Collection::stream)
+					.forEach(cookie -> getHeaders().add(headerName, cookie.toString()));
+		}
 		return this.headers;
 	}
 
+	@Override
 	public MultiValueMap<String, ResponseCookie> getCookies() {
 		return this.cookies;
 	}

spring-test/src/main/java/org/springframework/mock/http/server/reactive/MockServerHttpRequest.java

@@ -480,14 +480,16 @@ public class MockServerHttpRequest extends AbstractServerHttpRequest {
 
 		@Override
 		public MockServerHttpRequest body(Publisher<? extends DataBuffer> body) {
-			applyCookies();
+			applyCookiesIfNecessary();
 			return new MockServerHttpRequest(this.method, this.url, this.contextPath,
 					this.headers, this.cookies, this.remoteAddress, body);
 		}
 
-		private void applyCookies() {
-			this.cookies.values().stream().flatMap(Collection::stream)
-					.forEach(cookie -> this.headers.add(HttpHeaders.COOKIE, cookie.toString()));
+		private void applyCookiesIfNecessary() {
+			if (this.headers.get(HttpHeaders.COOKIE) == null) {
+				this.cookies.values().stream().flatMap(Collection::stream)
+						.forEach(cookie -> this.headers.add(HttpHeaders.COOKIE, cookie.toString()));
+			}
 		}
 	}
 

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

@@ -19,6 +19,7 @@ package org.springframework.test.web.reactive.server;
 import java.net.URI;
 import java.time.Duration;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.function.Function;
 
 import org.junit.Test;
@@ -28,6 +29,7 @@ import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DefaultDataBufferFactory;
 import org.springframework.core.io.buffer.support.DataBufferTestUtils;
 import org.springframework.http.HttpCookie;
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.ReactiveHttpOutputMessage;
@@ -69,8 +71,10 @@ public class HttpHandlerConnectorTests {
 		assertEquals(HttpMethod.POST, request.getMethod());
 		assertEquals("/custom-path", request.getURI().toString());
 
-		assertEquals(Arrays.asList("h0", "h1"), request.getHeaders().get("custom-header"));
+		HttpHeaders headers = request.getHeaders();
+		assertEquals(Arrays.asList("h0", "h1"), headers.get("custom-header"));
 		assertEquals(new HttpCookie("custom-cookie", "c0"), request.getCookies().getFirst("custom-cookie"));
+		assertEquals(Collections.singletonList("custom-cookie=c0"), headers.get(HttpHeaders.COOKIE));
 
 		DataBuffer buffer = request.getBody().blockFirst(Duration.ZERO);
 		assertEquals("Custom body", DataBufferTestUtils.dumpString(buffer, UTF_8));
@@ -93,8 +97,10 @@ public class HttpHandlerConnectorTests {
 				.block(Duration.ofSeconds(5));
 
 		assertEquals(HttpStatus.OK, response.getStatusCode());
-		assertEquals(Arrays.asList("h0", "h1"), response.getHeaders().get("custom-header"));
+		HttpHeaders headers = response.getHeaders();
+		assertEquals(Arrays.asList("h0", "h1"), headers.get("custom-header"));
 		assertEquals(cookie, response.getCookies().getFirst("custom-cookie"));
+		assertEquals(Collections.singletonList("custom-cookie=c0"), headers.get(HttpHeaders.SET_COOKIE));
 
 		DataBuffer buffer = response.getBody().blockFirst(Duration.ZERO);
 		assertEquals("Custom body", DataBufferTestUtils.dumpString(buffer, UTF_8));

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

@@ -15,45 +15,49 @@
  */
 package org.springframework.test.web.reactive.server;
 
+import java.util.Arrays;
+
 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.HttpHeaders;
+import org.springframework.http.HttpStatus;
 import org.springframework.http.ResponseCookie;
-import org.springframework.web.server.WebHandler;
+import org.springframework.http.server.reactive.ServerHttpResponse;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
 
 /**
- * Mock server integration test scenarios.
+ * Test scenarios involving a mock server.
  * @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.getAttributeOrDefault("foo", "none"))
-						.flatMap(value -> {
-							byte[] bytes = value.getBytes(UTF_8);
-							DataBuffer buffer = new DefaultDataBufferFactory().wrap(bytes);
-							return exchange.getResponse().writeWith(Mono.just(buffer));
-						});
-			}
-		};
-
-		WebTestClient client = new DefaultMockServerSpec(webHandler).configureClient().build();
+public class MockServerTests {
+
+
+	@Test // SPR-15674 (in comments)
+	public void mutateDoesNotCreateNewSession() throws Exception {
+
+		WebTestClient client = WebTestClient
+				.bindToWebHandler(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.getAttributeOrDefault("foo", "none"))
+								.flatMap(value -> {
+									byte[] bytes = value.getBytes(UTF_8);
+									DataBuffer buffer = new DefaultDataBufferFactory().wrap(bytes);
+									return exchange.getResponse().writeWith(Mono.just(buffer));
+								});
+					}
+				})
+				.build();
 
 		// Set the session attribute
 		EntityExchangeResult<Void> result = client.get().uri("/set").exchange()
@@ -69,10 +73,13 @@ public class MockServerIntegrationTests {
 				.expectBody(String.class).isEqualTo("bar");
 	}
 
-	@Test
+	@Test // SPR-16059
 	public void mutateDoesCopy() throws Exception {
 
-		WebTestClient.Builder builder = WebTestClient.bindToWebHandler(exchange -> exchange.getResponse().setComplete()).configureClient();
+		WebTestClient.Builder builder = WebTestClient
+				.bindToWebHandler(exchange -> exchange.getResponse().setComplete())
+				.configureClient();
+
 		builder.filter((request, next) -> next.exchange(request));
 		builder.defaultHeader("foo", "bar");
 		builder.defaultCookie("foo", "bar");
@@ -103,5 +110,30 @@ public class MockServerIntegrationTests {
 		clientFromMutatedBuilder.mutate().defaultCookies(cookies -> assertEquals(2, cookies.size()));
 	}
 
+	@Test // SPR-16124
+	public void exchangeResultHasCookieHeaders() throws Exception {
+
+		ExchangeResult result = WebTestClient
+				.bindToWebHandler(exchange -> {
+					ServerHttpResponse response = exchange.getResponse();
+					if (exchange.getRequest().getURI().getPath().equals("/cookie")) {
+						response.addCookie(ResponseCookie.from("a", "alpha").path("/pathA").build());
+						response.addCookie(ResponseCookie.from("b", "beta").path("/pathB").build());
+					}
+					else {
+						response.setStatusCode(HttpStatus.NOT_FOUND);
+					}
+					return response.setComplete();
+				})
+				.build()
+				.get().uri("/cookie").cookie("a", "alpha").cookie("b", "beta")
+				.exchange()
+				.expectStatus().isOk()
+				.expectHeader().valueEquals(HttpHeaders.SET_COOKIE, "a=alpha; Path=/pathA", "b=beta; Path=/pathB")
+				.expectBody().isEmpty();
+
+		assertEquals(Arrays.asList("a=alpha", "b=beta"),
+				result.getRequestHeaders().get(HttpHeaders.COOKIE));
+	}
 
 }

spring-web/src/test/java/org/springframework/mock/http/client/reactive/test/MockClientHttpRequest.java

@@ -17,6 +17,7 @@
 package org.springframework.mock.http.client.reactive.test;
 
 import java.net.URI;
+import java.util.Collection;
 import java.util.function.Function;
 
 import org.reactivestreams.Publisher;
@@ -26,6 +27,7 @@ import reactor.core.publisher.Mono;
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
 import org.springframework.core.io.buffer.DefaultDataBufferFactory;
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.client.reactive.AbstractClientHttpRequest;
 import org.springframework.http.client.reactive.ClientHttpRequest;
@@ -114,6 +116,8 @@ public class MockClientHttpRequest extends AbstractClientHttpRequest {
 
 	@Override
 	protected void applyCookies() {
+		getCookies().values().stream().flatMap(Collection::stream)
+				.forEach(cookie -> getHeaders().add(HttpHeaders.COOKIE, cookie.toString()));
 	}
 
 	@Override

spring-web/src/test/java/org/springframework/mock/http/client/reactive/test/MockClientHttpResponse.java

@@ -19,6 +19,7 @@ package org.springframework.mock.http.client.reactive.test;
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
+import java.util.Collection;
 
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
@@ -69,6 +70,11 @@ public class MockClientHttpResponse implements ClientHttpResponse {
 
 	@Override
 	public HttpHeaders getHeaders() {
+		String headerName = HttpHeaders.SET_COOKIE;
+		if (!getCookies().isEmpty() && this.headers.get(headerName) == null) {
+			getCookies().values().stream().flatMap(Collection::stream)
+					.forEach(cookie -> getHeaders().add(headerName, cookie.toString()));
+		}
 		return this.headers;
 	}
 
@@ -132,4 +138,4 @@ public class MockClientHttpResponse implements ClientHttpResponse {
 		return (charset != null ? charset : StandardCharsets.UTF_8);
 	}
 
-}
+}

spring-web/src/test/java/org/springframework/mock/http/server/reactive/test/MockServerHttpRequest.java

@@ -480,14 +480,16 @@ public class MockServerHttpRequest extends AbstractServerHttpRequest {
 
 		@Override
 		public MockServerHttpRequest body(Publisher<? extends DataBuffer> body) {
-			applyCookies();
+			applyCookiesIfNecessary();
 			return new MockServerHttpRequest(this.method, this.url, this.contextPath,
 					this.headers, this.cookies, this.remoteAddress, body);
 		}
 
-		private void applyCookies() {
-			this.cookies.values().stream().flatMap(Collection::stream)
-					.forEach(cookie -> this.headers.add(HttpHeaders.COOKIE, cookie.toString()));
+		private void applyCookiesIfNecessary() {
+			if (this.headers.get(HttpHeaders.COOKIE) == null) {
+				this.cookies.values().stream().flatMap(Collection::stream)
+						.forEach(cookie -> this.headers.add(HttpHeaders.COOKIE, cookie.toString()));
+			}
 		}
 	}