Allow to set cookies in ServerResponse

Issue: SPR-16121

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/DefaultEntityResponseBuilder.java

@@ -27,6 +27,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.Consumer;
 
 import reactor.core.publisher.Mono;
 
@@ -35,10 +36,13 @@ import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
+import org.springframework.http.ResponseCookie;
 import org.springframework.http.codec.HttpMessageWriter;
 import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.util.Assert;
+import org.springframework.util.LinkedMultiValueMap;
+import org.springframework.util.MultiValueMap;
 import org.springframework.web.reactive.function.BodyInserter;
 import org.springframework.web.server.ServerWebExchange;
 
@@ -58,6 +62,8 @@ class DefaultEntityResponseBuilder<T> implements EntityResponse.Builder<T> {
 
 	private final HttpHeaders headers = new HttpHeaders();
 
+	private final MultiValueMap<String, ResponseCookie> cookies = new LinkedMultiValueMap<>();
+
 	private final Map<String, Object> hints = new HashMap<>();
 
 
@@ -74,6 +80,21 @@ class DefaultEntityResponseBuilder<T> implements EntityResponse.Builder<T> {
 		return this;
 	}
 
+	@Override
+	public EntityResponse.Builder<T> cookie(ResponseCookie cookie) {
+		Assert.notNull(cookie, "'cookie' must not be null");
+		this.cookies.add(cookie.getName(), cookie);
+		return this;
+	}
+
+	@Override
+	public EntityResponse.Builder<T> cookies(
+			Consumer<MultiValueMap<String, ResponseCookie>> cookiesConsumer) {
+		Assert.notNull(cookiesConsumer, "'cookiesConsumer' must not be null");
+		cookiesConsumer.accept(this.cookies);
+		return this;
+	}
+
 	@Override
 	public EntityResponse.Builder<T> header(String headerName, String... headerValues) {
 		for (String headerValue : headerValues) {
@@ -161,8 +182,8 @@ class DefaultEntityResponseBuilder<T> implements EntityResponse.Builder<T> {
 
 	@Override
 	public Mono<EntityResponse<T>> build() {
-		return Mono.just(new DefaultEntityResponse<T>(this.status, this.headers, this.entity,
+		return Mono.just(new DefaultEntityResponse<T>(this.status, this.headers, this.cookies,
-				this.inserter, this.hints));
+				this.entity, this.inserter, this.hints));
 	}
 
 
@@ -177,10 +198,11 @@ class DefaultEntityResponseBuilder<T> implements EntityResponse.Builder<T> {
 		private final Map<String, Object> hints;
 
 
-		public DefaultEntityResponse(HttpStatus statusCode, HttpHeaders headers, T entity,
+		public DefaultEntityResponse(HttpStatus statusCode, HttpHeaders headers,
+				MultiValueMap<String, ResponseCookie> cookies, T entity,
 				BodyInserter<T, ? super ServerHttpResponse> inserter, Map<String, Object> hints) {
 
-			super(statusCode, headers);
+			super(statusCode, headers, cookies);
 			this.entity = entity;
 			this.inserter = inserter;
 			this.hints = hints;

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/DefaultRenderingResponseBuilder.java

@@ -22,6 +22,7 @@ import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.Locale;
 import java.util.Map;
+import java.util.function.Consumer;
 import java.util.stream.Stream;
 
 import reactor.core.publisher.Flux;
@@ -32,9 +33,12 @@ import org.springframework.core.Conventions;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
+import org.springframework.http.ResponseCookie;
 import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
+import org.springframework.util.LinkedMultiValueMap;
+import org.springframework.util.MultiValueMap;
 import org.springframework.web.reactive.result.view.ViewResolver;
 import org.springframework.web.server.ServerWebExchange;
 
@@ -52,6 +56,8 @@ class DefaultRenderingResponseBuilder implements RenderingResponse.Builder {
 
 	private final HttpHeaders headers = new HttpHeaders();
 
+	private final MultiValueMap<String, ResponseCookie> cookies = new LinkedMultiValueMap<>();
+
 	private final Map<String, Object> model = new LinkedHashMap<>();
 
 
@@ -67,6 +73,20 @@ class DefaultRenderingResponseBuilder implements RenderingResponse.Builder {
 		return this;
 	}
 
+	@Override
+	public RenderingResponse.Builder cookie(ResponseCookie cookie) {
+		Assert.notNull(cookie, "'cookie' must not be null");
+		this.cookies.add(cookie.getName(), cookie);
+		return this;
+	}
+
+	@Override
+	public RenderingResponse.Builder cookies(Consumer<MultiValueMap<String, ResponseCookie>> cookiesConsumer) {
+		Assert.notNull(cookiesConsumer, "'cookiesConsumer' must not be null");
+		cookiesConsumer.accept(this.cookies);
+		return this;
+	}
+
 	@Override
 	public RenderingResponse.Builder modelAttribute(Object attribute) {
 		Assert.notNull(attribute, "'value' must not be null");
@@ -117,7 +137,8 @@ class DefaultRenderingResponseBuilder implements RenderingResponse.Builder {
 
 	@Override
 	public Mono<RenderingResponse> build() {
-		return Mono.just(new DefaultRenderingResponse(this.status, this.headers, this.name, this.model));
+		return Mono.just(new DefaultRenderingResponse(this.status, this.headers, this.cookies,
+				this.name, this.model));
 	}
 
 
@@ -129,9 +150,10 @@ class DefaultRenderingResponseBuilder implements RenderingResponse.Builder {
 
 		private final Map<String, Object> model;
 
-		public DefaultRenderingResponse(HttpStatus statusCode, HttpHeaders headers, String name,
+		public DefaultRenderingResponse(HttpStatus statusCode, HttpHeaders headers,
-				Map<String, Object> model) {
+				MultiValueMap<String, ResponseCookie> cookies,
-			super(statusCode, headers);
+				String name, Map<String, Object> model) {
+			super(statusCode, headers, cookies);
 			this.name = name;
 			this.model = unmodifiableCopy(model);
 		}

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/DefaultServerResponseBuilder.java

@@ -37,10 +37,14 @@ import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
+import org.springframework.http.ResponseCookie;
 import org.springframework.http.codec.HttpMessageWriter;
 import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.util.Assert;
+import org.springframework.util.CollectionUtils;
+import org.springframework.util.LinkedMultiValueMap;
+import org.springframework.util.MultiValueMap;
 import org.springframework.web.reactive.function.BodyInserter;
 import org.springframework.web.reactive.function.BodyInserters;
 import org.springframework.web.server.ServerWebExchange;
@@ -57,6 +61,8 @@ class DefaultServerResponseBuilder implements ServerResponse.BodyBuilder {
 
 	private final HttpHeaders headers = new HttpHeaders();
 
+	private final MultiValueMap<String, ResponseCookie> cookies = new LinkedMultiValueMap<>();
+
 	private final Map<String, Object> hints = new HashMap<>();
 
 
@@ -80,6 +86,21 @@ class DefaultServerResponseBuilder implements ServerResponse.BodyBuilder {
 		return this;
 	}
 
+	@Override
+	public ServerResponse.BodyBuilder cookie(ResponseCookie cookie) {
+		Assert.notNull(cookie, "'cookie' must not be null");
+		this.cookies.add(cookie.getName(), cookie);
+		return this;
+	}
+
+	@Override
+	public ServerResponse.BodyBuilder cookies(
+			Consumer<MultiValueMap<String, ResponseCookie>> cookiesConsumer) {
+		Assert.notNull(cookiesConsumer, "'cookiesConsumer' must not be null");
+		cookiesConsumer.accept(this.cookies);
+		return this;
+	}
+
 	@Override
 	public ServerResponse.BodyBuilder allow(HttpMethod... allowedMethods) {
 		this.headers.setAllow(new LinkedHashSet<>(Arrays.asList(allowedMethods)));
@@ -166,7 +187,8 @@ class DefaultServerResponseBuilder implements ServerResponse.BodyBuilder {
 			BiFunction<ServerWebExchange, ServerResponse.Context, Mono<Void>> writeFunction) {
 
 		Assert.notNull(writeFunction, "'writeFunction' must not be null");
-		return Mono.just(new WriterFunctionServerResponse(this.statusCode, this.headers, writeFunction));
+		return Mono.just(new WriterFunctionServerResponse(this.statusCode, this.headers,
+				this.cookies, writeFunction));
 	}
 
 	@Override
@@ -214,7 +236,8 @@ class DefaultServerResponseBuilder implements ServerResponse.BodyBuilder {
 	@Override
 	public Mono<ServerResponse> body(BodyInserter<?, ? super ServerHttpResponse> inserter) {
 		Assert.notNull(inserter, "'inserter' must not be null");
-		return Mono.just(new BodyInserterServerResponse<>(this.statusCode, this.headers, inserter, this.hints));
+		return Mono.just(new BodyInserterServerResponse<>(this.statusCode, this.headers,
+				this.cookies, inserter, this.hints));
 	}
 
 	@Override
@@ -248,9 +271,13 @@ class DefaultServerResponseBuilder implements ServerResponse.BodyBuilder {
 
 		private final HttpHeaders headers;
 
-		protected AbstractServerResponse(HttpStatus statusCode, HttpHeaders headers) {
+		private final MultiValueMap<String, ResponseCookie> cookies;
+
+		protected AbstractServerResponse(HttpStatus statusCode, HttpHeaders headers,
+				MultiValueMap<String, ResponseCookie> cookies) {
 			this.statusCode = statusCode;
 			this.headers = readOnlyCopy(headers);
+			this.cookies = readOnlyCopy(cookies);
 		}
 
 		private static HttpHeaders readOnlyCopy(HttpHeaders headers) {
@@ -259,6 +286,12 @@ class DefaultServerResponseBuilder implements ServerResponse.BodyBuilder {
 			return HttpHeaders.readOnlyHttpHeaders(copy);
 		}
 
+		private static <K, V> MultiValueMap<K,V> readOnlyCopy(MultiValueMap<K, V> map) {
+			MultiValueMap<K, V> copy = new LinkedMultiValueMap<>();
+			copy.putAll(map);
+			return CollectionUtils.unmodifiableMultiValueMap(copy);
+		}
+
 		@Override
 		public final HttpStatus statusCode() {
 			return this.statusCode;
@@ -269,16 +302,23 @@ class DefaultServerResponseBuilder implements ServerResponse.BodyBuilder {
 			return this.headers;
 		}
 
+		@Override
+		public MultiValueMap<String, ResponseCookie> cookies() {
+			return this.cookies;
+		}
+
 		protected void writeStatusAndHeaders(ServerHttpResponse response) {
 			response.setStatusCode(this.statusCode);
-			HttpHeaders responseHeaders = response.getHeaders();
 
-			HttpHeaders headers = headers();
+			copy(this.headers, response.getHeaders());
-			if (!headers.isEmpty()) {
+			copy(this.cookies, response.getCookies());
-				headers.entrySet().stream()
+		}
-						.filter(entry -> !responseHeaders.containsKey(entry.getKey()))
+
-						.forEach(entry -> responseHeaders
+		private static <K,V> void copy(MultiValueMap<K,V> src, MultiValueMap<K,V> dst) {
-								.put(entry.getKey(), entry.getValue()));
+			if (!src.isEmpty()) {
+				src.entrySet().stream()
+						.filter(entry -> !dst.containsKey(entry.getKey()))
+						.forEach(entry -> dst.put(entry.getKey(), entry.getValue()));
 			}
 		}
 	}
@@ -289,9 +329,10 @@ class DefaultServerResponseBuilder implements ServerResponse.BodyBuilder {
 		private final BiFunction<ServerWebExchange, Context, Mono<Void>> writeFunction;
 
 		public WriterFunctionServerResponse(HttpStatus statusCode, HttpHeaders headers,
+				MultiValueMap<String, ResponseCookie> cookies,
 				BiFunction<ServerWebExchange, Context, Mono<Void>> writeFunction) {
 
-			super(statusCode, headers);
+			super(statusCode, headers, cookies);
 			this.writeFunction = writeFunction;
 		}
 
@@ -310,9 +351,10 @@ class DefaultServerResponseBuilder implements ServerResponse.BodyBuilder {
 		private final Map<String, Object> hints;
 
 		public BodyInserterServerResponse(HttpStatus statusCode, HttpHeaders headers,
+				MultiValueMap<String, ResponseCookie> cookies,
 				BodyInserter<T, ? super ServerHttpResponse> inserter, Map<String, Object> hints) {
 
-			super(statusCode, headers);
+			super(statusCode, headers, cookies);
 			this.inserter = inserter;
 			this.hints = hints;
 		}

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/EntityResponse.java

@@ -19,6 +19,7 @@ package org.springframework.web.reactive.function.server;
 import java.net.URI;
 import java.time.ZonedDateTime;
 import java.util.Set;
+import java.util.function.Consumer;
 
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Mono;
@@ -29,8 +30,10 @@ import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
+import org.springframework.http.ResponseCookie;
 import org.springframework.http.codec.json.Jackson2CodecSupport;
 import org.springframework.http.server.reactive.ServerHttpResponse;
+import org.springframework.util.MultiValueMap;
 import org.springframework.web.reactive.function.BodyInserter;
 import org.springframework.web.reactive.function.BodyInserters;
 
@@ -123,6 +126,24 @@ public interface EntityResponse<T> extends ServerResponse {
 		 */
 		Builder<T> status(HttpStatus status);
 
+		/**
+		 * Add the given cookie to the response.
+		 * @param cookie the cookie to add
+		 * @return this builder
+		 */
+		Builder<T> cookie(ResponseCookie cookie);
+
+		/**
+		 * Manipulate this response's cookies with the given consumer. The
+		 * cookies provided to the consumer are "live", so that the consumer can be used to
+		 * {@linkplain MultiValueMap#set(Object, Object) overwrite} existing cookies,
+		 * {@linkplain MultiValueMap#remove(Object) remove} cookies, or use any of the other
+		 * {@link MultiValueMap} methods.
+		 * @param cookiesConsumer a function that consumes the cookies
+		 * @return this builder
+		 */
+		Builder<T> cookies(Consumer<MultiValueMap<String, ResponseCookie>> cookiesConsumer);
+
 		/**
 		 * Set the set of allowed {@link HttpMethod HTTP methods}, as specified
 		 * by the {@code Allow} header.

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/RenderingResponse.java

@@ -18,13 +18,16 @@ package org.springframework.web.reactive.function.server;
 
 import java.util.Collection;
 import java.util.Map;
+import java.util.function.Consumer;
 
 import reactor.core.publisher.Mono;
 
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpStatus;
+import org.springframework.http.ResponseCookie;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
+import org.springframework.util.MultiValueMap;
 
 /**
  * Rendering-specific subtype of {@link ServerResponse} that exposes model and template data.
@@ -139,6 +142,24 @@ public interface RenderingResponse extends ServerResponse {
 		 */
 		Builder status(HttpStatus status);
 
+		/**
+		 * Add the given cookie to the response.
+		 * @param cookie the cookie to add
+		 * @return this builder
+		 */
+		Builder cookie(ResponseCookie cookie);
+
+		/**
+		 * Manipulate this response's cookies with the given consumer. The
+		 * cookies provided to the consumer are "live", so that the consumer can be used to
+		 * {@linkplain MultiValueMap#set(Object, Object) overwrite} existing cookies,
+		 * {@linkplain MultiValueMap#remove(Object) remove} cookies, or use any of the other
+		 * {@link MultiValueMap} methods.
+		 * @param cookiesConsumer a function that consumes the cookies
+		 * @return this builder
+		 */
+		Builder cookies(Consumer<MultiValueMap<String, ResponseCookie>> cookiesConsumer);
+
 		/**
 		 * Build the response.
 		 * @return the built response

spring-webflux/src/main/java/org/springframework/web/reactive/function/server/ServerResponse.java

@@ -34,10 +34,12 @@ import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
+import org.springframework.http.ResponseCookie;
 import org.springframework.http.codec.HttpMessageWriter;
 import org.springframework.http.codec.json.Jackson2CodecSupport;
 import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.util.Assert;
+import org.springframework.util.MultiValueMap;
 import org.springframework.web.reactive.function.BodyInserter;
 import org.springframework.web.reactive.function.BodyInserters;
 import org.springframework.web.reactive.result.view.ViewResolver;
@@ -66,6 +68,11 @@ public interface ServerResponse {
 	 */
 	HttpHeaders headers();
 
+	/**
+	 * Return the cookies of this response.
+	 */
+	MultiValueMap<String, ResponseCookie> cookies();
+
 	/**
 	 * Write this response to the given web exchange.
 	 * @param exchange the web exchange to write to
@@ -219,6 +226,24 @@ public interface ServerResponse {
 		 */
 		B headers(Consumer<HttpHeaders> headersConsumer);
 
+		/**
+		 * Add the given cookie to the response.
+		 * @param cookie the cookie to add
+		 * @return this builder
+		 */
+		B cookie(ResponseCookie cookie);
+
+		/**
+		 * Manipulate this response's cookies with the given consumer. The
+		 * cookies provided to the consumer are "live", so that the consumer can be used to
+		 * {@linkplain MultiValueMap#set(Object, Object) overwrite} existing cookies,
+		 * {@linkplain MultiValueMap#remove(Object) remove} cookies, or use any of the other
+		 * {@link MultiValueMap} methods.
+		 * @param cookiesConsumer a function that consumes the cookies
+		 * @return this builder
+		 */
+		B cookies(Consumer<MultiValueMap<String, ResponseCookie>> cookiesConsumer);
+
 		/**
 		 * Set the set of allowed {@link HttpMethod HTTP methods}, as specified
 		 * by the {@code Allow} header.

spring-webflux/src/test/java/org/springframework/web/reactive/function/server/DefaultEntityResponseBuilderTests.java

@@ -39,17 +39,19 @@ import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
+import org.springframework.http.ResponseCookie;
 import org.springframework.http.codec.EncoderHttpMessageWriter;
 import org.springframework.http.codec.HttpMessageWriter;
 import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.mock.http.server.reactive.test.MockServerHttpRequest;
 import org.springframework.mock.web.test.server.MockServerWebExchange;
+import org.springframework.util.LinkedMultiValueMap;
+import org.springframework.util.MultiValueMap;
 import org.springframework.web.reactive.function.BodyInserter;
 import org.springframework.web.reactive.result.view.ViewResolver;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
-import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.*;
-import static org.junit.Assert.assertSame;
 
 /**
  * @author Arjen Poutsma
@@ -176,6 +178,18 @@ public class DefaultEntityResponseBuilderTests {
 				.verify();
 	}
 
+	@Test
+	public void cookies() throws Exception {
+		MultiValueMap<String, ResponseCookie> newCookies = new LinkedMultiValueMap<>();
+		newCookies.add("name", ResponseCookie.from("name", "value").build());
+		Mono<EntityResponse<String>> result =
+				EntityResponse.fromObject("foo").cookies(cookies -> cookies.addAll(newCookies)).build();
+		StepVerifier.create(result)
+				.expectNextMatches(response -> newCookies.equals(response.cookies()))
+				.expectComplete()
+				.verify();
+	}
+
 	@Test
 	public void bodyInserter() throws Exception {
 		String body = "foo";

spring-webflux/src/test/java/org/springframework/web/reactive/function/server/DefaultRenderingResponseTests.java

@@ -28,8 +28,11 @@ import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
 import org.springframework.http.HttpHeaders;
+import org.springframework.http.ResponseCookie;
 import org.springframework.mock.http.server.reactive.test.MockServerHttpRequest;
 import org.springframework.mock.web.test.server.MockServerWebExchange;
+import org.springframework.util.LinkedMultiValueMap;
+import org.springframework.util.MultiValueMap;
 import org.springframework.web.reactive.result.view.View;
 import org.springframework.web.reactive.result.view.ViewResolver;
 
@@ -103,6 +106,19 @@ public class DefaultRenderingResponseTests {
 				.verify();
 	}
 
+	@Test
+	public void cookies() throws Exception {
+		MultiValueMap<String, ResponseCookie> newCookies = new LinkedMultiValueMap<>();
+		newCookies.add("name", ResponseCookie.from("name", "value").build());
+		Mono<RenderingResponse> result =
+				RenderingResponse.create("foo").cookies(cookies -> cookies.addAll(newCookies)).build();
+		StepVerifier.create(result)
+				.expectNextMatches(response -> newCookies.equals(response.cookies()))
+				.expectComplete()
+				.verify();
+	}
+
+
 	@Test
 	public void render() throws Exception {
 		Map<String, Object> model = Collections.singletonMap("foo", "bar");

spring-webflux/src/test/java/org/springframework/web/reactive/function/server/DefaultServerResponseBuilderTests.java

@@ -32,7 +32,10 @@ import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
+import org.springframework.http.ResponseCookie;
 import org.springframework.mock.http.server.reactive.test.MockServerHttpResponse;
+import org.springframework.util.LinkedMultiValueMap;
+import org.springframework.util.MultiValueMap;
 import org.springframework.web.server.ServerWebExchange;
 
 import static org.junit.Assert.*;
@@ -264,10 +267,25 @@ public class DefaultServerResponseBuilderTests {
 
 	}
 
+	@Test
+	public void cookies() throws Exception {
+		MultiValueMap<String, ResponseCookie> newCookies = new LinkedMultiValueMap<>();
+		newCookies.add("name", ResponseCookie.from("name", "value").build());
+		Mono<ServerResponse> result =
+				ServerResponse.ok().cookies(cookies -> cookies.addAll(newCookies)).build();
+		StepVerifier.create(result)
+				.expectNextMatches(response -> newCookies.equals(response.cookies()))
+				.expectComplete()
+				.verify();
+	}
+
 	@Test
 	public void build() throws Exception {
+		ResponseCookie cookie = ResponseCookie.from("name", "value").build();
 		Mono<ServerResponse>
-				result = ServerResponse.status(HttpStatus.CREATED).header("MyKey", "MyValue").build();
+				result = ServerResponse.status(HttpStatus.CREATED)
+				.header("MyKey", "MyValue")
+				.cookie(cookie).build();
 
 		ServerWebExchange exchange = mock(ServerWebExchange.class);
 		MockServerHttpResponse response = new MockServerHttpResponse();
@@ -278,6 +296,7 @@ public class DefaultServerResponseBuilderTests {
 
 		assertEquals(HttpStatus.CREATED, response.getStatusCode());
 		assertEquals("MyValue", response.getHeaders().getFirst("MyKey"));
+		assertEquals("value", response.getCookies().getFirst("name").getValue());
 		StepVerifier.create(response.getBody()).expectComplete().verify();
 	}
 

spring-webflux/src/test/java/org/springframework/web/reactive/function/server/RouterFunctionsTests.java

@@ -25,9 +25,12 @@ import reactor.test.StepVerifier;
 
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpStatus;
+import org.springframework.http.ResponseCookie;
 import org.springframework.http.server.reactive.HttpHandler;
 import org.springframework.mock.http.server.reactive.test.MockServerHttpRequest;
 import org.springframework.mock.http.server.reactive.test.MockServerHttpResponse;
+import org.springframework.util.LinkedMultiValueMap;
+import org.springframework.util.MultiValueMap;
 import org.springframework.web.server.ResponseStatusException;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebFilter;
@@ -197,6 +200,11 @@ public class RouterFunctionsTests {
 						return new HttpHeaders();
 					}
 
+					@Override
+					public MultiValueMap<String, ResponseCookie> cookies() {
+						return new LinkedMultiValueMap<>();
+					}
+
 					@Override
 					public Mono<Void> writeTo(ServerWebExchange exchange,
 							Context context) {
@@ -230,6 +238,11 @@ public class RouterFunctionsTests {
 						return new HttpHeaders();
 					}
 
+					@Override
+					public MultiValueMap<String, ResponseCookie> cookies() {
+						return new LinkedMultiValueMap<>();
+					}
+
 					@Override
 					public Mono<Void> writeTo(ServerWebExchange exchange,
 							Context context) {