diff --git a/spring-test/src/main/java/org/springframework/mock/web/MockHttpServletResponse.java b/spring-test/src/main/java/org/springframework/mock/web/MockHttpServletResponse.java index f9b179d60e..018689ced0 100644 --- a/spring-test/src/main/java/org/springframework/mock/web/MockHttpServletResponse.java +++ b/spring-test/src/main/java/org/springframework/mock/web/MockHttpServletResponse.java @@ -55,6 +55,7 @@ import org.springframework.web.util.WebUtils; * @author Rod Johnson * @author Brian Clozel * @author Vedran Pavic + * @author Sam Brannen * @since 1.0.2 */ public class MockHttpServletResponse implements HttpServletResponse { @@ -573,20 +574,22 @@ public class MockHttpServletResponse implements HttpServletResponse { } private void setHeaderValue(String name, Object value) { - if (setSpecialHeader(name, value)) { + boolean replaceHeader = true; + if (setSpecialHeader(name, value, replaceHeader)) { return; } - doAddHeaderValue(name, value, true); + doAddHeaderValue(name, value, replaceHeader); } private void addHeaderValue(String name, Object value) { - if (setSpecialHeader(name, value)) { + boolean replaceHeader = false; + if (setSpecialHeader(name, value, replaceHeader)) { return; } - doAddHeaderValue(name, value, false); + doAddHeaderValue(name, value, replaceHeader); } - private boolean setSpecialHeader(String name, Object value) { + private boolean setSpecialHeader(String name, Object value, boolean replaceHeader) { if (HttpHeaders.CONTENT_TYPE.equalsIgnoreCase(name)) { setContentType(value.toString()); return true; @@ -605,7 +608,12 @@ public class MockHttpServletResponse implements HttpServletResponse { } else if (HttpHeaders.SET_COOKIE.equalsIgnoreCase(name)) { MockCookie cookie = MockCookie.parse(value.toString()); - addCookie(cookie); + if (replaceHeader) { + setCookie(cookie); + } + else { + addCookie(cookie); + } return true; } else { @@ -628,6 +636,20 @@ public class MockHttpServletResponse implements HttpServletResponse { } } + /** + * Set the {@code Set-Cookie} header to the supplied {@link Cookie}, + * overwriting any previous cookies. + * @param cookie the {@code Cookie} to set + * @since 5.1.10 + * @see #addCookie(Cookie) + */ + private void setCookie(Cookie cookie) { + Assert.notNull(cookie, "Cookie must not be null"); + this.cookies.clear(); + this.cookies.add(cookie); + doAddHeaderValue(HttpHeaders.SET_COOKIE, getCookieHeader(cookie), true); + } + @Override public void setStatus(int status) { if (!this.isCommitted()) { diff --git a/spring-test/src/test/java/org/springframework/mock/web/MockHttpServletResponseTests.java b/spring-test/src/test/java/org/springframework/mock/web/MockHttpServletResponseTests.java index 708e2e1549..a5e32599d9 100644 --- a/spring-test/src/test/java/org/springframework/mock/web/MockHttpServletResponseTests.java +++ b/spring-test/src/test/java/org/springframework/mock/web/MockHttpServletResponseTests.java @@ -322,22 +322,36 @@ public class MockHttpServletResponseTests { assertEquals(HttpServletResponse.SC_NOT_FOUND, response.getStatus()); } + /** + * @since 5.1.10 + */ @Test - public void setCookieHeaderValid() { + public void setCookieHeader() { + response.setHeader(HttpHeaders.SET_COOKIE, "SESSION=123; Path=/; Secure; HttpOnly; SameSite=Lax"); + assertNumCookies(1); + assertPrimarySessionCookie("123"); + + // Setting the Set-Cookie header a 2nd time should overwrite the previous value + response.setHeader(HttpHeaders.SET_COOKIE, "SESSION=999; Path=/; Secure; HttpOnly; SameSite=Lax"); + assertNumCookies(1); + assertPrimarySessionCookie("999"); + } + + @Test + public void addCookieHeader() { response.addHeader(HttpHeaders.SET_COOKIE, "SESSION=123; Path=/; Secure; HttpOnly; SameSite=Lax"); - Cookie cookie = response.getCookie("SESSION"); - assertNotNull(cookie); - assertTrue(cookie instanceof MockCookie); - assertEquals("SESSION", cookie.getName()); - assertEquals("123", cookie.getValue()); - assertEquals("/", cookie.getPath()); - assertTrue(cookie.getSecure()); - assertTrue(cookie.isHttpOnly()); - assertEquals("Lax", ((MockCookie) cookie).getSameSite()); + assertNumCookies(1); + assertPrimarySessionCookie("123"); + + // Adding a 2nd cookie header should result in 2 cookies. + response.addHeader(HttpHeaders.SET_COOKIE, "SESSION=999; Path=/; Secure; HttpOnly; SameSite=Lax"); + assertNumCookies(2); + assertPrimarySessionCookie("123"); + assertCookieValues("123", "999"); } @Test - public void addMockCookie() { + public void addCookie() { MockCookie mockCookie = new MockCookie("SESSION", "123"); mockCookie.setPath("/"); mockCookie.setDomain("example.com"); @@ -348,9 +362,35 @@ public class MockHttpServletResponseTests { response.addCookie(mockCookie); + assertNumCookies(1); assertEquals("SESSION=123; Path=/; Domain=example.com; Max-Age=0; " + "Expires=Thu, 01 Jan 1970 00:00:00 GMT; Secure; HttpOnly; SameSite=Lax", response.getHeader(HttpHeaders.SET_COOKIE)); + + // Adding a 2nd Cookie should result in 2 Cookies. + response.addCookie(new MockCookie("SESSION", "999")); + assertNumCookies(2); + assertCookieValues("123", "999"); + } + + private void assertNumCookies(int expected) { + assertEquals(expected, this.response.getCookies().length); + } + + private void assertCookieValues(String... expected) { + assertArrayEquals(expected, Arrays.stream(response.getCookies()).map(Cookie::getValue).toArray()); + } + + private void assertPrimarySessionCookie(String expectedValue) { + Cookie cookie = this.response.getCookie("SESSION"); + assertNotNull(cookie); + assertTrue(cookie instanceof MockCookie); + assertEquals("SESSION", cookie.getName()); + assertEquals(expectedValue, cookie.getValue()); + assertEquals("/", cookie.getPath()); + assertTrue(cookie.getSecure()); + assertTrue(cookie.isHttpOnly()); + assertEquals("Lax", ((MockCookie) cookie).getSameSite()); } } diff --git a/spring-web/src/test/java/org/springframework/mock/web/test/MockHttpServletResponse.java b/spring-web/src/test/java/org/springframework/mock/web/test/MockHttpServletResponse.java index 50861b3a72..ece0503b30 100644 --- a/spring-web/src/test/java/org/springframework/mock/web/test/MockHttpServletResponse.java +++ b/spring-web/src/test/java/org/springframework/mock/web/test/MockHttpServletResponse.java @@ -55,6 +55,7 @@ import org.springframework.web.util.WebUtils; * @author Rod Johnson * @author Brian Clozel * @author Vedran Pavic + * @author Sam Brannen * @since 1.0.2 */ public class MockHttpServletResponse implements HttpServletResponse { @@ -573,20 +574,22 @@ public class MockHttpServletResponse implements HttpServletResponse { } private void setHeaderValue(String name, Object value) { - if (setSpecialHeader(name, value)) { + boolean replaceHeader = true; + if (setSpecialHeader(name, value, replaceHeader)) { return; } - doAddHeaderValue(name, value, true); + doAddHeaderValue(name, value, replaceHeader); } private void addHeaderValue(String name, Object value) { - if (setSpecialHeader(name, value)) { + boolean replaceHeader = false; + if (setSpecialHeader(name, value, replaceHeader)) { return; } - doAddHeaderValue(name, value, false); + doAddHeaderValue(name, value, replaceHeader); } - private boolean setSpecialHeader(String name, Object value) { + private boolean setSpecialHeader(String name, Object value, boolean replaceHeader) { if (HttpHeaders.CONTENT_TYPE.equalsIgnoreCase(name)) { setContentType(value.toString()); return true; @@ -605,7 +608,12 @@ public class MockHttpServletResponse implements HttpServletResponse { } else if (HttpHeaders.SET_COOKIE.equalsIgnoreCase(name)) { MockCookie cookie = MockCookie.parse(value.toString()); - addCookie(cookie); + if (replaceHeader) { + setCookie(cookie); + } + else { + addCookie(cookie); + } return true; } else { @@ -628,6 +636,20 @@ public class MockHttpServletResponse implements HttpServletResponse { } } + /** + * Set the {@code Set-Cookie} header to the supplied {@link Cookie}, + * overwriting any previous cookies. + * @param cookie the {@code Cookie} to set + * @since 5.1.10 + * @see #addCookie(Cookie) + */ + private void setCookie(Cookie cookie) { + Assert.notNull(cookie, "Cookie must not be null"); + this.cookies.clear(); + this.cookies.add(cookie); + doAddHeaderValue(HttpHeaders.SET_COOKIE, getCookieHeader(cookie), true); + } + @Override public void setStatus(int status) { if (!this.isCommitted()) {