From e86becc151d655f09f02f7eed6d31cc34a75b49e Mon Sep 17 00:00:00 2001
From: Rob Winch <rwinch@users.noreply.github.com>
Date: Thu, 15 Mar 2018 08:28:47 -0500
Subject: [PATCH] Relax assertions in HeaderSpecTests

Fixes: gh-5116
---
 .../config/web/server/HeaderSpecTests.java    | 33 +++++++++++--------
 1 file changed, 19 insertions(+), 14 deletions(-)

diff --git a/config/src/test/java/org/springframework/security/config/web/server/HeaderSpecTests.java b/config/src/test/java/org/springframework/security/config/web/server/HeaderSpecTests.java
index 6cb2c65e45..faafc01871 100644
--- a/config/src/test/java/org/springframework/security/config/web/server/HeaderSpecTests.java
+++ b/config/src/test/java/org/springframework/security/config/web/server/HeaderSpecTests.java
@@ -28,7 +28,7 @@ import org.springframework.test.web.reactive.server.FluxExchangeResult;
 import org.springframework.test.web.reactive.server.WebTestClient;
 
 import java.time.Duration;
-import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -45,7 +45,7 @@ public class HeaderSpecTests {
 
 	HttpHeaders expectedHeaders = new HttpHeaders();
 
-	Set<String> ignoredHeaderNames = Collections.singleton(HttpHeaders.CONTENT_TYPE);
+	Set<String> headerNamesNotPresent = new HashSet<>();
 
 	@Before
 	public void setup() {
@@ -67,9 +67,7 @@ public class HeaderSpecTests {
 
 	@Test
 	public void headersWhenCacheDisableThenCacheNotWritten() {
-		this.expectedHeaders.remove(HttpHeaders.CACHE_CONTROL);
-		this.expectedHeaders.remove(HttpHeaders.PRAGMA);
-		this.expectedHeaders.remove(HttpHeaders.EXPIRES);
+		expectHeaderNamesNotPresent(HttpHeaders.CACHE_CONTROL, HttpHeaders.PRAGMA, HttpHeaders.EXPIRES);
 		this.headers.cache().disable();
 
 		assertHeaders();
@@ -77,7 +75,7 @@ public class HeaderSpecTests {
 
 	@Test
 	public void headersWhenContentOptionsDisableThenContentTypeOptionsNotWritten() {
-		this.expectedHeaders.remove(ContentTypeOptionsServerHttpHeadersWriter.X_CONTENT_OPTIONS);
+		expectHeaderNamesNotPresent(ContentTypeOptionsServerHttpHeadersWriter.X_CONTENT_OPTIONS);
 		this.headers.contentTypeOptions().disable();
 
 		assertHeaders();
@@ -85,7 +83,7 @@ public class HeaderSpecTests {
 
 	@Test
 	public void headersWhenHstsDisableThenHstsNotWritten() {
-		this.expectedHeaders.remove(StrictTransportSecurityServerHttpHeadersWriter.STRICT_TRANSPORT_SECURITY);
+		expectHeaderNamesNotPresent(StrictTransportSecurityServerHttpHeadersWriter.STRICT_TRANSPORT_SECURITY);
 		this.headers.hsts().disable();
 
 		assertHeaders();
@@ -103,7 +101,7 @@ public class HeaderSpecTests {
 
 	@Test
 	public void headersWhenFrameOptionsDisableThenFrameOptionsNotWritten() {
-		this.expectedHeaders.remove(XFrameOptionsServerHttpHeadersWriter.X_FRAME_OPTIONS);
+		expectHeaderNamesNotPresent(XFrameOptionsServerHttpHeadersWriter.X_FRAME_OPTIONS);
 		this.headers.frameOptions().disable();
 
 		assertHeaders();
@@ -111,9 +109,7 @@ public class HeaderSpecTests {
 
 	@Test
 	public void headersWhenFrameOptionsModeThenFrameOptionsCustomMode() {
-		this.expectedHeaders.remove(XFrameOptionsServerHttpHeadersWriter.X_FRAME_OPTIONS);
-		this.expectedHeaders
-			.add(XFrameOptionsServerHttpHeadersWriter.X_FRAME_OPTIONS, "SAMEORIGIN");
+		this.expectedHeaders.set(XFrameOptionsServerHttpHeadersWriter.X_FRAME_OPTIONS, "SAMEORIGIN");
 		this.headers.frameOptions().mode(XFrameOptionsServerHttpHeadersWriter.Mode.SAMEORIGIN);
 
 		assertHeaders();
@@ -121,12 +117,19 @@ public class HeaderSpecTests {
 
 	@Test
 	public void headersWhenXssProtectionDisableThenXssProtectionNotWritten() {
-		this.expectedHeaders.remove("X-Xss-Protection");
+		expectHeaderNamesNotPresent("X-Xss-Protection");
 		this.headers.xssProtection().disable();
 
 		assertHeaders();
 	}
 
+	private void expectHeaderNamesNotPresent(String... headerNames) {
+		for(String headerName : headerNames) {
+			this.expectedHeaders.remove(headerName);
+			this.headerNamesNotPresent.add(headerName);
+		}
+	}
+
 	private void assertHeaders() {
 		WebTestClient client = buildClient();
 		FluxExchangeResult<String> response = client.get()
@@ -135,10 +138,12 @@ public class HeaderSpecTests {
 			.returnResult(String.class);
 
 		Map<String, List<String>> responseHeaders = response.getResponseHeaders();
-		this.ignoredHeaderNames.stream().forEach(responseHeaders::remove);
 
-		assertThat(responseHeaders).describedAs(response.toString()).isEqualTo(
+		assertThat(responseHeaders).describedAs(response.toString()).containsAllEntriesOf(
 			this.expectedHeaders);
+		if (!this.headerNamesNotPresent.isEmpty()) {
+			assertThat(responseHeaders.keySet()).doesNotContainAnyElementsOf(this.headerNamesNotPresent);
+		}
 	}
 
 	private WebTestClient buildClient() {