MockMvcWebConnection stores cookies from response

Previously MockMvcWebConnection did not update the cookie manager with the
cookies from MockHttpServletResponse. This meant that newly added cookies
are not saved to the cookie manager and thus are not presented in the next
request.

This commit ensures that MockMvcWebConnection stores the response cookies
in the cookie manager.

Issue: SPR-14265

Author: Rob Winch

spring-test/src/main/java/org/springframework/test/web/servlet/htmlunit/MockMvcWebConnection.java

@@ -17,13 +17,16 @@
 package org.springframework.test.web.servlet.htmlunit;
 
 import java.io.IOException;
+import java.util.Date;
 import java.util.HashMap;
 import java.util.Map;
 
+import com.gargoylesoftware.htmlunit.CookieManager;
 import com.gargoylesoftware.htmlunit.WebClient;
 import com.gargoylesoftware.htmlunit.WebConnection;
 import com.gargoylesoftware.htmlunit.WebRequest;
 import com.gargoylesoftware.htmlunit.WebResponse;
+import com.gargoylesoftware.htmlunit.util.Cookie;
 
 import org.springframework.mock.web.MockHttpServletResponse;
 import org.springframework.mock.web.MockHttpSession;
@@ -142,10 +145,32 @@ public WebResponse getResponse(WebRequest webRequest) throws IOException {
 			httpServletResponse = getResponse(requestBuilder);
 			forwardedUrl = httpServletResponse.getForwardedUrl();
 		}
+		storeCookies(webRequest, httpServletResponse.getCookies());
 
 		return new MockWebResponseBuilder(startTime, webRequest, httpServletResponse).build();
 	}
 
+	private void storeCookies(WebRequest webRequest, javax.servlet.http.Cookie[] cookies) {
+		if (cookies == null) {
+			return;
+		}
+		Date now = new Date();
+		CookieManager cookieManager = webClient.getCookieManager();
+		for (javax.servlet.http.Cookie cookie : cookies) {
+			if (cookie.getDomain() == null) {
+				cookie.setDomain(webRequest.getUrl().getHost());
+			}
+			Cookie toManage = MockWebResponseBuilder.createCookie(cookie);
+			Date expires = toManage.getExpires();
+			if (expires == null || expires.after(now)) {
+				cookieManager.addCookie(toManage);
+			}
+			else {
+				cookieManager.removeCookie(toManage);
+			}
+		}
+	}
+
 	private MockHttpServletResponse getResponse(RequestBuilder requestBuilder) throws IOException {
 		ResultActions resultActions;
 		try {
@@ -182,5 +207,4 @@ static void validateContextPath(String contextPath) {
 			throw new IllegalArgumentException("contextPath '" + contextPath + "' must not end with '/'.");
 		}
 	}
-
 }

spring-test/src/main/java/org/springframework/test/web/servlet/htmlunit/MockWebResponseBuilder.java

@@ -112,6 +112,10 @@ private List<NameValuePair> responseHeaders() {
 	}
 
 	private String valueOfCookie(Cookie cookie) {
+		return createCookie(cookie).toString();
+	}
+
+	static com.gargoylesoftware.htmlunit.util.Cookie createCookie(Cookie cookie) {
 		Date expires = null;
 		if (cookie.getMaxAge() > -1) {
 			expires = new Date(System.currentTimeMillis() + cookie.getMaxAge() * 1000);
@@ -125,7 +129,6 @@ private String valueOfCookie(Cookie cookie) {
 		if(cookie.isHttpOnly()) {
 			result.setAttribute("httponly", "true");
 		}
-		return new com.gargoylesoftware.htmlunit.util.Cookie(result).toString();
+		return new com.gargoylesoftware.htmlunit.util.Cookie(result);
 	}
-
 }

spring-test/src/test/java/org/springframework/test/web/servlet/htmlunit/MockMvcWebClientBuilderTests.java

@@ -19,7 +19,9 @@
 import java.io.IOException;
 import java.net.URL;
 import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletResponse;
 
+import com.gargoylesoftware.htmlunit.HttpMethod;
 import com.gargoylesoftware.htmlunit.WebClient;
 import com.gargoylesoftware.htmlunit.WebRequest;
 import com.gargoylesoftware.htmlunit.WebResponse;
@@ -38,7 +40,10 @@
 import org.springframework.tests.Assume;
 import org.springframework.tests.TestGroup;
 import org.springframework.web.bind.annotation.CookieValue;
+import org.springframework.web.bind.annotation.DeleteMapping;
+import org.springframework.web.bind.annotation.PostMapping;
 import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.RestController;
 import org.springframework.web.context.WebApplicationContext;
 import org.springframework.web.servlet.config.annotation.EnableWebMvc;
@@ -103,11 +108,22 @@ public void cookieManagerShared() throws Exception {
 		this.mockMvc = MockMvcBuilders.standaloneSetup(new CookieController()).build();
 		WebClient client = MockMvcWebClientBuilder.mockMvcSetup(this.mockMvc).build();
 
-		assertThat(getResponse(client, "http://localhost/").getContentAsString(), equalTo(""));
+		assertThat(getResponse(client, "http://localhost/").getContentAsString(), equalTo("NA"));
 		client.getCookieManager().addCookie(new Cookie("localhost", "cookie", "cookieManagerShared"));
 		assertThat(getResponse(client, "http://localhost/").getContentAsString(), equalTo("cookieManagerShared"));
 	}
 
+	@Test // SPR-14265
+	public void cookiesAreManaged() throws Exception {
+		this.mockMvc = MockMvcBuilders.standaloneSetup(new CookieController()).build();
+		WebClient client = MockMvcWebClientBuilder.mockMvcSetup(this.mockMvc).build();
+
+		assertThat(getResponse(client, "http://localhost/").getContentAsString(), equalTo("NA"));
+		assertThat(postResponse(client, "http://localhost/?cookie=foo").getContentAsString(), equalTo("Set"));
+		assertThat(getResponse(client, "http://localhost/").getContentAsString(), equalTo("foo"));
+		assertThat(deleteResponse(client, "http://localhost/").getContentAsString(), equalTo("Delete"));
+		assertThat(getResponse(client, "http://localhost/").getContentAsString(), equalTo("NA"));
+	}
 
 	private void assertMockMvcUsed(WebClient client, String url) throws Exception {
 		assertThat(getResponse(client, url).getContentAsString(), equalTo("mvc"));
@@ -118,7 +134,19 @@ private void assertMockMvcNotUsed(WebClient client, String url) throws Exception
 	}
 
 	private WebResponse getResponse(WebClient client, String url) throws IOException {
-		return client.getWebConnection().getResponse(new WebRequest(new URL(url)));
+		return createResponse(client, new WebRequest(new URL(url)));
+	}
+
+	private WebResponse postResponse(WebClient client, String url) throws IOException {
+		return createResponse(client, new WebRequest(new URL(url), HttpMethod.POST));
+	}
+
+	private WebResponse deleteResponse(WebClient client, String url) throws IOException {
+		return createResponse(client, new WebRequest(new URL(url), HttpMethod.DELETE));
+	}
+
+	private WebResponse createResponse(WebClient client, WebRequest request) throws IOException {
+		return client.getWebConnection().getResponse(request);
 	}
 
 
@@ -139,10 +167,26 @@ public String contextPath(HttpServletRequest request) {
 	@RestController
 	static class CookieController {
 
+		static final String COOKIE_NAME = "cookie";
+
 		@RequestMapping(path = "/", produces = "text/plain")
-		String cookie(@CookieValue("cookie") String cookie) {
+		String cookie(@CookieValue(name = COOKIE_NAME, defaultValue = "NA") String cookie) {
 			return cookie;
 		}
+
+		@PostMapping(path = "/", produces = "text/plain")
+		String setCookie(@RequestParam String cookie, HttpServletResponse response) {
+			response.addCookie(new javax.servlet.http.Cookie(COOKIE_NAME, cookie));
+			return "Set";
+		}
+
+		@DeleteMapping(path = "/", produces = "text/plain")
+		String deleteCookie(HttpServletResponse response) {
+			javax.servlet.http.Cookie cookie = new javax.servlet.http.Cookie(COOKIE_NAME, "");
+			cookie.setMaxAge(0);
+			response.addCookie(cookie);
+			return "Delete";
+		}
 	}
 
 }