From 0d42a1bd7fdd44b000873edd6cbbb35697fddc4a Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@pivotal.io>
Date: Thu, 19 Mar 2020 14:32:56 +0000
Subject: [PATCH] Add retry for flaky test (suspected Tomcat issue)

---
 .../AbstractWebSocketIntegrationTests.java     |  7 +++----
 .../socket/WebSocketIntegrationTests.java      | 18 +++++++++++++++---
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/spring-webflux/src/test/java/org/springframework/web/reactive/socket/AbstractWebSocketIntegrationTests.java b/spring-webflux/src/test/java/org/springframework/web/reactive/socket/AbstractWebSocketIntegrationTests.java
index 02f8a6593de..8d0e78cf2a4 100644
--- a/spring-webflux/src/test/java/org/springframework/web/reactive/socket/AbstractWebSocketIntegrationTests.java
+++ b/spring-webflux/src/test/java/org/springframework/web/reactive/socket/AbstractWebSocketIntegrationTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -23,7 +23,6 @@ import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 import java.net.URI;
-import java.net.URISyntaxException;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.stream.Stream;
@@ -156,8 +155,8 @@ abstract class AbstractWebSocketIntegrationTests {
 		return WebHttpHandlerBuilder.applicationContext(context).build();
 	}
 
-	protected URI getUrl(String path) throws URISyntaxException {
-		return new URI("ws://localhost:" + this.port + path);
+	protected URI getUrl(String path) {
+		return URI.create("ws://localhost:" + this.port + path);
 	}
 
 	protected abstract Class<?> getWebConfigClass();
diff --git a/spring-webflux/src/test/java/org/springframework/web/reactive/socket/WebSocketIntegrationTests.java b/spring-webflux/src/test/java/org/springframework/web/reactive/socket/WebSocketIntegrationTests.java
index 4ac544d83e8..9be00e9d353 100644
--- a/spring-webflux/src/test/java/org/springframework/web/reactive/socket/WebSocketIntegrationTests.java
+++ b/spring-webflux/src/test/java/org/springframework/web/reactive/socket/WebSocketIntegrationTests.java
@@ -29,6 +29,7 @@ import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.MonoProcessor;
 import reactor.core.publisher.ReplayProcessor;
+import reactor.util.retry.Retry;
 
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -39,6 +40,7 @@ import org.springframework.web.reactive.handler.SimpleUrlHandlerMapping;
 import org.springframework.web.reactive.socket.client.WebSocketClient;
 import org.springframework.web.server.WebFilter;
 import org.springframework.web.testfixture.http.server.reactive.bootstrap.HttpServer;
+import org.springframework.web.testfixture.http.server.reactive.bootstrap.TomcatHttpServer;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -66,18 +68,28 @@ class WebSocketIntegrationTests extends AbstractWebSocketIntegrationTests {
 	void echo(WebSocketClient client, HttpServer server, Class<?> serverConfigClass) throws Exception {
 		startServer(client, server, serverConfigClass);
 
+		if (server instanceof TomcatHttpServer) {
+			Mono.fromRunnable(this::testEcho)
+					.retryWhen(Retry.max(3).filter(ex -> ex instanceof IllegalStateException))
+					.block();
+		}
+		else {
+			testEcho();
+		}
+	}
+
+	private void testEcho() {
 		int count = 100;
 		Flux<String> input = Flux.range(1, count).map(index -> "msg-" + index);
 		ReplayProcessor<Object> output = ReplayProcessor.create(count);
-
 		this.client.execute(getUrl("/echo"), session -> session
 				.send(input.map(session::textMessage))
 				.thenMany(session.receive().take(count).map(WebSocketMessage::getPayloadAsText))
 				.subscribeWith(output)
 				.then())
 				.block(TIMEOUT);
-
-		assertThat(output.collectList().block(TIMEOUT)).isEqualTo(input.collectList().block(TIMEOUT));
+		assertThat(output.isTerminated()).isTrue();
+		assertThat(output.collectList().block()).isEqualTo(input.collectList().block());
 	}
 
 	@ParameterizedWebSocketTest