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 4121d9a8b5..69e31d51a4 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 @@ -57,6 +57,7 @@ import org.springframework.web.util.WebUtils; * @author Brian Clozel * @author Vedran Pavic * @author Sebastien Deleuze + * @author Sam Brannen * @since 1.0.2 */ public class MockHttpServletResponse implements HttpServletResponse { @@ -602,20 +603,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; @@ -634,7 +637,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 { @@ -657,6 +665,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 f1ff821139..4be36a65dc 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 @@ -333,23 +333,36 @@ class MockHttpServletResponseTests { assertThat(response.getStatus()).isEqualTo(HttpServletResponse.SC_NOT_FOUND); } + /** + * @since 5.1.10 + */ @Test - void setCookieHeaderValid() { - response.addHeader(HttpHeaders.SET_COOKIE, "SESSION=123; Path=/; Secure; HttpOnly; SameSite=Lax"); - Cookie cookie = response.getCookie("SESSION"); - assertThat(cookie).isNotNull(); - boolean condition = cookie instanceof MockCookie; - assertThat(condition).isTrue(); - assertThat(cookie.getName()).isEqualTo("SESSION"); - assertThat(cookie.getValue()).isEqualTo("123"); - assertThat(cookie.getPath()).isEqualTo("/"); - assertThat(cookie.getSecure()).isTrue(); - assertThat(cookie.isHttpOnly()).isTrue(); - assertThat(((MockCookie) cookie).getSameSite()).isEqualTo("Lax"); + 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 - void addMockCookie() { + void addCookieHeader() { + response.addHeader(HttpHeaders.SET_COOKIE, "SESSION=123; Path=/; Secure; HttpOnly; SameSite=Lax"); + 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 + void addCookie() { MockCookie mockCookie = new MockCookie("SESSION", "123"); mockCookie.setPath("/"); mockCookie.setDomain("example.com"); @@ -360,8 +373,35 @@ class MockHttpServletResponseTests { response.addCookie(mockCookie); + assertNumCookies(1); assertThat(response.getHeader(HttpHeaders.SET_COOKIE)).isEqualTo(("SESSION=123; Path=/; Domain=example.com; Max-Age=0; " + "Expires=Thu, 01 Jan 1970 00:00:00 GMT; Secure; HttpOnly; SameSite=Lax")); + + // 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) { + assertThat(this.response.getCookies()).hasSize(expected); + } + + private void assertCookieValues(String... expected) { + assertThat(response.getCookies()).extracting(Cookie::getValue).containsExactly(expected); + } + + private void assertPrimarySessionCookie(String expectedValue) { + Cookie cookie = this.response.getCookie("SESSION"); + assertThat(cookie).isNotNull(); + boolean condition = cookie instanceof MockCookie; + assertThat(condition).isTrue(); + assertThat(cookie.getName()).isEqualTo("SESSION"); + assertThat(cookie.getValue()).isEqualTo(expectedValue); + assertThat(cookie.getPath()).isEqualTo("/"); + assertThat(cookie.getSecure()).isTrue(); + assertThat(cookie.isHttpOnly()).isTrue(); + assertThat(((MockCookie) cookie).getSameSite()).isEqualTo("Lax"); } } 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 97a39464f4..3f992773b8 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 @@ -57,6 +57,7 @@ import org.springframework.web.util.WebUtils; * @author Brian Clozel * @author Vedran Pavic * @author Sebastien Deleuze + * @author Sam Brannen * @since 1.0.2 */ public class MockHttpServletResponse implements HttpServletResponse { @@ -602,20 +603,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; @@ -634,7 +637,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 { @@ -657,6 +665,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()) {