Include cookies in request snippets

spring-restdocs-core/src/main/java/org/springframework/restdocs/cli/CliOperationRequest.java

@@ -25,6 +25,8 @@
 import java.util.Map.Entry;
 import java.util.Set;
 
+import javax.servlet.http.Cookie;
+
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.restdocs.operation.OperationRequest;
@@ -125,6 +127,11 @@ public URI getUri() {
 		return this.delegate.getUri();
 	}
 
+	@Override
+	public Collection<Cookie> getCookies() {
+		return this.delegate.getCookies();
+	}
+
 	private interface HeaderFilter {
 
 		boolean allow(String name, List<String> value);

spring-restdocs-core/src/main/java/org/springframework/restdocs/cli/CurlRequestSnippet.java

@@ -23,6 +23,8 @@
 import java.util.Map;
 import java.util.Map.Entry;
 
+import javax.servlet.http.Cookie;
+
 import org.springframework.http.HttpMethod;
 import org.springframework.restdocs.operation.Operation;
 import org.springframework.restdocs.operation.OperationRequest;
@@ -92,12 +94,28 @@ private String getOptions(Operation operation) {
 		writeUserOptionIfNecessary(request, printer);
 		writeHttpMethodIfNecessary(request, printer);
 		writeHeaders(request, printer);
+		writeCookies(request, printer);
 		writePartsIfNecessary(request, printer);
 		writeContent(request, printer);
 
 		return command.toString();
 	}
 
+	private void writeCookies(CliOperationRequest request, PrintWriter printer) {
+		if (request.getCookies() != null && request.getCookies().size() > 0) {
+			printer.print(" --cookie ");
+			StringBuilder cookiesBuilder = new StringBuilder();
+
+			for (Cookie cookie : request.getCookies()) {
+				cookiesBuilder.append(String.format("%s=%s;", cookie.getName(), cookie.getValue()));
+			}
+
+			String cookiesHeader = cookiesBuilder.substring(0, cookiesBuilder.length() - 1); // remove trailing semicolon
+
+			printer.print(String.format("'%s'", cookiesHeader)); // add single quotes
+		}
+	}
+
 	private void writeIncludeHeadersInOutputOption(PrintWriter writer) {
 		writer.print("-i");
 	}

spring-restdocs-core/src/main/java/org/springframework/restdocs/cli/HttpieRequestSnippet.java

@@ -23,6 +23,8 @@
 import java.util.Map;
 import java.util.Map.Entry;
 
+import javax.servlet.http.Cookie;
+
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.MediaType;
@@ -160,6 +162,10 @@ private void writeHeaders(OperationRequest request, PrintWriter writer) {
 				writer.print(String.format(" '%s:%s'", entry.getKey(), header));
 			}
 		}
+
+		for (Cookie cookie : request.getCookies()) {
+			writer.print(String.format(" 'Cookie:%s=%s'", cookie.getName(), cookie.getValue()));
+		}
 	}
 
 	private void writeParametersIfNecessary(CliOperationRequest request,

spring-restdocs-core/src/main/java/org/springframework/restdocs/http/HttpRequestSnippet.java

@@ -24,6 +24,8 @@
 import java.util.Map;
 import java.util.Map.Entry;
 
+import javax.servlet.http.Cookie;
+
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.MediaType;
@@ -35,6 +37,7 @@
 import org.springframework.restdocs.snippet.TemplatedSnippet;
 import org.springframework.util.StringUtils;
 
+
 /**
  * A {@link Snippet} that documents an HTTP request.
  *
@@ -112,6 +115,11 @@ private List<Map<String, String>> getHeaders(OperationRequest request) {
 
 			}
 		}
+
+		for (Cookie cookie : request.getCookies()) {
+			headers.add(header(HttpHeaders.COOKIE, String.format("%s=%s", cookie.getName(), cookie.getValue())));
+		}
+
 		if (requiresFormEncodingContentTypeHeader(request)) {
 			headers.add(header(HttpHeaders.CONTENT_TYPE,
 					MediaType.APPLICATION_FORM_URLENCODED_VALUE));

spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/OperationRequest.java

@@ -19,6 +19,8 @@
 import java.net.URI;
 import java.util.Collection;
 
+import javax.servlet.http.Cookie;
+
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 
@@ -86,4 +88,11 @@ public interface OperationRequest {
 	 */
 	URI getUri();
 
+	/**
+	 * Returns cookies sent with the request.
+	 *
+	 * @return the cookies
+	 */
+	Collection<Cookie> getCookies();
+
 }

spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/OperationRequestFactory.java

@@ -18,6 +18,9 @@
 
 import java.net.URI;
 import java.util.Collection;
+import java.util.Collections;
+
+import javax.servlet.http.Cookie;
 
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
@@ -40,13 +43,34 @@ public class OperationRequestFactory {
 	 * @param headers the request's headers
 	 * @param parameters the request's parameters
 	 * @param parts the request's parts
+	 * @param cookies the request's cookies
 	 * @return the {@code OperationRequest}
 	 */
 	public OperationRequest create(URI uri, HttpMethod method, byte[] content,
 			HttpHeaders headers, Parameters parameters,
-			Collection<OperationRequestPart> parts) {
+			Collection<OperationRequestPart> parts,
+			Collection<Cookie> cookies) {
 		return new StandardOperationRequest(uri, method, content,
-				augmentHeaders(headers, uri, content), parameters, parts);
+				augmentHeaders(headers, uri, content), parameters, parts, cookies);
+	}
+
+	/**
+	 * Creates a new {@link OperationRequest}. The given {@code headers} will be augmented
+	 * to ensure that they always include a {@code Content-Length} header if the request
+	 * has any content and a {@code Host} header.
+	 *
+	 * @param uri the request's uri
+	 * @param method the request method
+	 * @param content the content of the request
+	 * @param headers the request's headers
+	 * @param parameters the request's parameters
+	 * @param parts the request's parts
+	 * @return the {@code OperationRequest}
+	 */
+	public OperationRequest create(URI uri, HttpMethod method, byte[] content,
+			HttpHeaders headers, Parameters parameters,
+			Collection<OperationRequestPart> parts) {
+		return create(uri, method, content, headers, parameters, parts, Collections.<Cookie>emptyList());
 	}
 
 	/**
@@ -62,7 +86,7 @@ public OperationRequest create(URI uri, HttpMethod method, byte[] content,
 	public OperationRequest createFrom(OperationRequest original, byte[] newContent) {
 		return new StandardOperationRequest(original.getUri(), original.getMethod(),
 				newContent, getUpdatedHeaders(original.getHeaders(), newContent),
-				original.getParameters(), original.getParts());
+				original.getParameters(), original.getParts(), original.getCookies());
 	}
 
 	/**
@@ -78,7 +102,7 @@ public OperationRequest createFrom(OperationRequest original,
 			HttpHeaders newHeaders) {
 		return new StandardOperationRequest(original.getUri(), original.getMethod(),
 				original.getContent(), newHeaders, original.getParameters(),
-				original.getParts());
+				original.getParts(), original.getCookies());
 	}
 
 	/**
@@ -94,7 +118,7 @@ public OperationRequest createFrom(OperationRequest original,
 			Parameters newParameters) {
 		return new StandardOperationRequest(original.getUri(), original.getMethod(),
 				original.getContent(), original.getHeaders(), newParameters,
-				original.getParts());
+				original.getParts(), original.getCookies());
 	}
 
 	private HttpHeaders augmentHeaders(HttpHeaders originalHeaders, URI uri,

spring-restdocs-core/src/main/java/org/springframework/restdocs/operation/StandardOperationRequest.java

@@ -20,6 +20,8 @@
 import java.util.Collection;
 import java.util.Collections;
 
+import javax.servlet.http.Cookie;
+
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 
@@ -39,6 +41,8 @@ class StandardOperationRequest extends AbstractOperationMessage
 
 	private URI uri;
 
+	private Collection<Cookie> cookies;
+
 	/**
 	 * Creates a new request with the given {@code uri} and {@code method}. The request
 	 * will have the given {@code headers}, {@code parameters}, and {@code parts}.
@@ -52,12 +56,14 @@ class StandardOperationRequest extends AbstractOperationMessage
 	 */
 	StandardOperationRequest(URI uri, HttpMethod method, byte[] content,
 			HttpHeaders headers, Parameters parameters,
-			Collection<OperationRequestPart> parts) {
+			Collection<OperationRequestPart> parts,
+			Collection<Cookie> cookies) {
 		super(content, headers);
 		this.uri = uri;
 		this.method = method;
 		this.parameters = parameters;
 		this.parts = parts;
+		this.cookies = cookies;
 	}
 
 	@Override
@@ -80,4 +86,9 @@ public URI getUri() {
 		return this.uri;
 	}
 
+	@Override
+	public Collection<Cookie> getCookies() {
+		return this.cookies;
+	}
+
 }

spring-restdocs-core/src/test/java/org/springframework/restdocs/cli/CurlRequestSnippetTests.java

@@ -251,6 +251,17 @@ public void requestWithHeaders() throws IOException {
 						.header("a", "alpha").build());
 	}
 
+	@Test
+	public void requestWithCookies() throws IOException {
+		this.snippets.expectCurlRequest()
+				.withContents(codeBlock("bash").content("$ curl 'http://localhost/foo' -i" +
+						" --cookie 'name1=value1;name2=value2'"));
+		new CurlRequestSnippet()
+				.document(this.operationBuilder.request("http://localhost/foo")
+						.cookie("name1", "value1")
+						.cookie("name2", "value2").build());
+	}
+
 	@Test
 	public void multipartPostWithNoSubmittedFileName() throws IOException {
 		String expectedContent = "$ curl 'http://localhost/upload' -i -X POST -H "

spring-restdocs-core/src/test/java/org/springframework/restdocs/cli/HttpieRequestSnippetTests.java

@@ -252,6 +252,17 @@ public void requestWithHeaders() throws IOException {
 						.header("a", "alpha").build());
 	}
 
+	@Test
+	public void requestWithCookies() throws IOException {
+		this.snippets.expectHttpieRequest().withContents(
+				codeBlock("bash").content("$ http GET 'http://localhost/foo'" +
+						" 'Cookie:name1=value1' 'Cookie:name2=value2'"));
+		new HttpieRequestSnippet()
+				.document(this.operationBuilder.request("http://localhost/foo")
+						.cookie("name1", "value1")
+						.cookie("name2", "value2").build());
+	}
+
 	@Test
 	public void multipartPostWithNoSubmittedFileName() throws IOException {
 		String expectedContent = String

spring-restdocs-core/src/test/java/org/springframework/restdocs/headers/RequestHeadersSnippetTests.java

@@ -54,7 +54,7 @@ public void requestWithHeaders() throws IOException {
 						.row("`X-Test`", "one").row("`Accept`", "two")
 						.row("`Accept-Encoding`", "three")
 						.row("`Accept-Language`", "four").row("`Cache-Control`", "five")
-						.row("`Connection`", "six"));
+						.row("`Connection`", "six").row("`Cookie`", "seven"));
 		new RequestHeadersSnippet(
 				Arrays.asList(headerWithName("X-Test").description("one"),
 						headerWithName("Accept").description("two"),
@@ -63,7 +63,8 @@ public void requestWithHeaders() throws IOException {
 						headerWithName("Cache-Control").description("five"),
 						headerWithName(
 								"Connection")
-										.description("six")))
+										.description("six"),
+						headerWithName("Cookie").description("seven")))
 												.document(
 														this.operationBuilder
 																.request(
@@ -78,6 +79,8 @@ public void requestWithHeaders() throws IOException {
 																		"max-age=0")
 																.header("Connection",
 																		"keep-alive")
+																.header("Cookie",
+																		"cookie1=cookieVal1; cookie2=cookieVal2")
 																.build());
 	}
 

spring-restdocs-core/src/test/java/org/springframework/restdocs/http/HttpRequestSnippetTests.java

@@ -80,6 +80,21 @@ public void getRequestWithPort() throws IOException {
 				.request("http://localhost:8080/foo").header("Alpha", "a").build());
 	}
 
+	@Test
+	public void getRequestWithCookies() throws IOException {
+		this.snippets.expectHttpRequest()
+				.withContents(httpRequest(RequestMethod.GET, "/foo")
+						.header(HttpHeaders.HOST, "localhost")
+						.header(HttpHeaders.COOKIE, "name1=value1")
+						.header(HttpHeaders.COOKIE, "name2=value2"));
+
+		new HttpRequestSnippet()
+				.document(this.operationBuilder.request("http://localhost/foo")
+				.cookie("name1", "value1")
+				.cookie("name2", "value2")
+				.build());
+	}
+
 	@Test
 	public void getRequestWithQueryString() throws IOException {
 		this.snippets.expectHttpRequest()

spring-restdocs-core/src/test/java/org/springframework/restdocs/test/OperationBuilder.java

@@ -19,11 +19,14 @@
 import java.io.File;
 import java.net.URI;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
+import javax.servlet.http.Cookie;
+
 import org.junit.runners.model.Statement;
 
 import org.springframework.http.HttpHeaders;
@@ -153,6 +156,8 @@ public final class OperationRequestBuilder {
 
 		private List<OperationRequestPartBuilder> partBuilders = new ArrayList<>();
 
+		private Collection<Cookie> cookies = new ArrayList<>();
+
 		private OperationRequestBuilder(String uri) {
 			this.requestUri = URI.create(uri);
 		}
@@ -163,7 +168,7 @@ private OperationRequest buildRequest() {
 				parts.add(builder.buildPart());
 			}
 			return new OperationRequestFactory().create(this.requestUri, this.method,
-					this.content, this.headers, this.parameters, parts);
+					this.content, this.headers, this.parameters, parts, this.cookies);
 		}
 
 		public Operation build() {
@@ -209,6 +214,11 @@ public OperationRequestPartBuilder part(String name, byte[] content) {
 			return partBuilder;
 		}
 
+		public OperationRequestBuilder cookie(String name, String value) {
+			this.cookies.add(new Cookie(name, value));
+			return this;
+		}
+
 		/**
 		 * Basic builder API for creating an {@link OperationRequestPart}.
 		 */