Improve the fluency of the ResponseCreator API

spring-test/src/main/java/org/springframework/test/web/client/response/DefaultResponseCreator.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2021 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.
@@ -19,18 +19,25 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URI;
+import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
+import java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import org.springframework.core.io.Resource;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
+import org.springframework.http.ResponseCookie;
 import org.springframework.http.client.ClientHttpRequest;
 import org.springframework.http.client.ClientHttpResponse;
 import org.springframework.lang.Nullable;
 import org.springframework.mock.http.client.MockClientHttpResponse;
 import org.springframework.test.web.client.ResponseCreator;
 import org.springframework.util.Assert;
+import org.springframework.util.MultiValueMap;
 
 /**
  * A {@code ResponseCreator} with builder-style methods for adding response details.
@@ -40,7 +47,7 @@
  */
 public class DefaultResponseCreator implements ResponseCreator {
 
-	private HttpStatus statusCode;
+	private final HttpStatus statusCode;
 
 	private byte[] content = new byte[0];
 
@@ -59,7 +66,6 @@ protected DefaultResponseCreator(HttpStatus statusCode) {
 		this.statusCode = statusCode;
 	}
 
-
 	/**
 	 * Set the body as a UTF-8 String.
 	 */
@@ -68,6 +74,14 @@ public DefaultResponseCreator body(String content) {
 		return this;
 	}
 
+	/**
+	 * Set the body from a string using the given character set.
+	 */
+	public DefaultResponseCreator body(String content, Charset charset) {
+		this.content = content.getBytes(charset);
+		return this;
+	}
+
 	/**
 	 * Set the body as a byte array.
 	 */
@@ -77,7 +91,7 @@ public DefaultResponseCreator body(byte[] content) {
 	}
 
 	/**
-	 * Set the body as a {@link Resource}.
+	 * Set the body from a {@link Resource}.
 	 */
 	public DefaultResponseCreator body(Resource resource) {
 		this.contentResource = resource;
@@ -100,6 +114,26 @@ public DefaultResponseCreator location(URI location) {
 		return this;
 	}
 
+	/**
+	 * Add a single header.
+	 */
+	public DefaultResponseCreator header(String name, String value) {
+		// This is really just an alias, but it makes the interface more fluent.
+		return headers(name, value);
+	}
+
+	/**
+	 * Add one or more headers.
+	 */
+	public DefaultResponseCreator headers(String name, String ... value) {
+		List<String> valueList = Stream.of(value)
+				.filter(Objects::nonNull)
+				.collect(Collectors.toList());
+
+		this.headers.addAll(name, valueList);
+		return this;
+	}
+
 	/**
 	 * Copy all given headers.
 	 */
@@ -108,6 +142,36 @@ public DefaultResponseCreator headers(HttpHeaders headers) {
 		return this;
 	}
 
+	/**
+	 * Add a single cookie.
+	 */
+	public DefaultResponseCreator cookie(ResponseCookie cookie) {
+		// This is really just an alias, but it makes the interface more fluent.
+		return cookies(cookie);
+	}
+
+	/**
+	 * Add one or more cookies.
+	 */
+	public DefaultResponseCreator cookies(ResponseCookie... cookies) {
+		for (ResponseCookie cookie : cookies) {
+			this.headers.add(HttpHeaders.SET_COOKIE, cookie.toString());
+		}
+
+		return this;
+	}
+
+	/**
+	 * Copy all given cookies.
+	 */
+	public DefaultResponseCreator cookies(MultiValueMap<String, ResponseCookie> cookies) {
+		cookies.values()
+				.stream()
+				.flatMap(List::stream)
+				.forEach(cookie -> this.headers.add(HttpHeaders.SET_COOKIE, cookie.toString()));
+
+		return this;
+	}
 
 	@Override
 	public ClientHttpResponse createResponse(@Nullable ClientHttpRequest request) throws IOException {

spring-test/src/main/java/org/springframework/test/web/client/response/MockRestResponseCreators.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2021 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.
@@ -20,6 +20,7 @@
 import java.net.URI;
 
 import org.springframework.core.io.Resource;
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.lang.Nullable;
@@ -81,6 +82,13 @@ public static DefaultResponseCreator withCreatedEntity(URI location) {
 		return new DefaultResponseCreator(HttpStatus.CREATED).location(location);
 	}
 
+	/**
+	 * {@code ResponseCreator} for a 202 response (ACCEPTED).
+	 */
+	public static DefaultResponseCreator withAccepted() {
+		return new DefaultResponseCreator(HttpStatus.ACCEPTED);
+	}
+
 	/**
 	 * {@code ResponseCreator} for a 204 response (NO_CONTENT).
 	 */
@@ -102,13 +110,71 @@ public static DefaultResponseCreator withUnauthorizedRequest() {
 		return new DefaultResponseCreator(HttpStatus.UNAUTHORIZED);
 	}
 
+	/**
+	 * {@code ResponseCreator} for a 403 response (FORBIDDEN).
+	 */
+	public static DefaultResponseCreator withForbiddenRequest() {
+		return new DefaultResponseCreator(HttpStatus.FORBIDDEN);
+	}
+
+	/**
+	 * {@code ResponseCreator} for a 404 response (NOT_FOUND).
+	 */
+	public static DefaultResponseCreator withResourceNotFound() {
+		return new DefaultResponseCreator(HttpStatus.NOT_FOUND);
+	}
+
+	/**
+	 * {@code ResponseCreator} for a 409 response (CONFLICT).
+	 */
+	public static DefaultResponseCreator withRequestConflict() {
+		return new DefaultResponseCreator(HttpStatus.CONFLICT);
+	}
+
+	/**
+	 * {@code ResponseCreator} for a 429 ratelimited response (TOO_MANY_REQUESTS).
+	 */
+	public static DefaultResponseCreator withTooManyRequests() {
+		return new DefaultResponseCreator(HttpStatus.TOO_MANY_REQUESTS);
+	}
+
+	/**
+	 * {@code ResponseCreator} for a 429 ratelimited response (TOO_MANY_REQUESTS) with a {@code Retry-After} header
+	 * in seconds.
+	 */
+	public static DefaultResponseCreator withTooManyRequests(int retryAfter) {
+		return new DefaultResponseCreator(HttpStatus.TOO_MANY_REQUESTS)
+				.header(HttpHeaders.RETRY_AFTER, Integer.toString(retryAfter));
+	}
+
 	/**
 	 * {@code ResponseCreator} for a 500 response (SERVER_ERROR).
 	 */
 	public static DefaultResponseCreator withServerError() {
 		return new DefaultResponseCreator(HttpStatus.INTERNAL_SERVER_ERROR);
 	}
 
+	/**
+	 * {@code ResponseCreator} for a 502 response (BAD_GATEWAY).
+	 */
+	public static DefaultResponseCreator withBadGateway() {
+		return new DefaultResponseCreator(HttpStatus.BAD_GATEWAY);
+	}
+
+	/**
+	 * {@code ResponseCreator} for a 503 response (SERVICE_UNAVAILABLE).
+	 */
+	public static DefaultResponseCreator withServiceUnavailable() {
+		return new DefaultResponseCreator(HttpStatus.SERVICE_UNAVAILABLE);
+	}
+
+	/**
+	 * {@code ResponseCreator} for a 504 response (GATEWAY_TIMEOUT).
+	 */
+	public static DefaultResponseCreator withGatewayTimeout() {
+		return new DefaultResponseCreator(HttpStatus.GATEWAY_TIMEOUT);
+	}
+
 	/**
 	 * {@code ResponseCreator} with a specific HTTP status.
 	 * @param status the response status