From 34a4a10277009fab36f7075b963659673b039c36 Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <marcin@grzejszczak.pl>
Date: Wed, 27 Feb 2019 15:24:02 +0100
Subject: [PATCH] Trying to make tests less brittle

---
 .../instrument/web/TraceWebFluxTests.java     | 34 +++++++++----------
 .../client/integration/WebClientTests.java    |  7 ----
 2 files changed, 17 insertions(+), 24 deletions(-)

diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFluxTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFluxTests.java
index 03b8c890f..57c13aa10 100644
--- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFluxTests.java
+++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFluxTests.java
@@ -440,16 +440,20 @@ class SleuthSpanCreatorAspectWebFlux {
 		Long newSpanId = (Long) object;
 
 		Awaitility.await().untilAsserted(() -> {
-			List<zipkin2.Span> spans = getSpans();
-			then(spans).hasSize(2);
-			then(spans.get(0).name()).isEqualTo("new-span-in-trace-context");
-			then(spans.get(0).id()).isEqualTo(toHexString(newSpanId));
-			then(spans.get(1).kind()).isEqualTo(zipkin2.Span.Kind.SERVER);
-			then(spans.get(1).name()).isEqualTo("get /test/newspan1");
+			then(spanWithName("new-span-in-trace-context").id())
+					.isEqualTo(toHexString(newSpanId));
+			then(spanWithName("get /test/newspan1").kind())
+					.isEqualTo(zipkin2.Span.Kind.SERVER);
 			then(this.tracer.currentSpan()).isNull();
 		});
 	}
 
+	private zipkin2.Span spanWithName(String name) {
+		return getSpans().stream().filter(span -> name.equals(span.name())).findFirst()
+				.orElseThrow(() -> new AssertionError(
+						"Span with name [" + name + "] not found"));
+	}
+
 	public void shouldCreateNewSpanInWebFluxInSubscriberContext() {
 		setup();
 		Mono<Object> mono = this.webClient.get().uri("/test/newSpan2").exchange()
@@ -460,12 +464,10 @@ class SleuthSpanCreatorAspectWebFlux {
 		Long newSpanId = (Long) object;
 
 		Awaitility.await().untilAsserted(() -> {
-			List<zipkin2.Span> spans = getSpans();
-			then(spans).hasSize(2);
-			then(spans.get(0).name()).isEqualTo("new-span-in-subscriber-context");
-			then(spans.get(0).id()).isEqualTo(toHexString(newSpanId));
-			then(spans.get(1).kind()).isEqualTo(zipkin2.Span.Kind.SERVER);
-			then(spans.get(1).name()).isEqualTo("get /test/newspan2");
+			then(spanWithName("new-span-in-subscriber-context").id())
+					.isEqualTo(toHexString(newSpanId));
+			then(spanWithName("get /test/newspan2").kind())
+					.isEqualTo(zipkin2.Span.Kind.SERVER);
 			then(this.tracer.currentSpan()).isNull();
 		});
 	}
@@ -481,12 +483,10 @@ class SleuthSpanCreatorAspectWebFlux {
 		Long newSpanId = (Long) object;
 
 		Awaitility.await().untilAsserted(() -> {
-			List<zipkin2.Span> spans = getSpans();
-			then(spans).hasSize(1);
-			then(spans.get(0).kind()).isEqualTo(zipkin2.Span.Kind.SERVER);
-			then(spans.get(0).name()).isEqualTo("get /test/ping");
+			then(spanWithName("get /test/ping").kind())
+					.isEqualTo(zipkin2.Span.Kind.SERVER);
 			then(this.repository.getSpan()).isNotNull();
-			then(spans.get(0).id()).isEqualTo(toHexString(newSpanId))
+			then(spanWithName("get /test/ping").id()).isEqualTo(toHexString(newSpanId))
 					.isEqualTo(this.repository.getSpan().context().traceIdString());
 			then(this.tracer.currentSpan()).isNull();
 		});
diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/integration/WebClientTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/integration/WebClientTests.java
index 4b07e8277..05853670a 100644
--- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/integration/WebClientTests.java
+++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/integration/WebClientTests.java
@@ -52,7 +52,6 @@ import org.apache.http.impl.nio.client.CloseableHttpAsyncClient;
 import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
 import org.awaitility.Awaitility;
 import org.junit.After;
-import org.junit.BeforeClass;
 import org.junit.ClassRule;
 import org.junit.Rule;
 import org.junit.Test;
@@ -78,7 +77,6 @@ import org.springframework.cloud.gateway.config.GatewayClassPathWarningAutoConfi
 import org.springframework.cloud.netflix.ribbon.RibbonClient;
 import org.springframework.cloud.openfeign.EnableFeignClients;
 import org.springframework.cloud.openfeign.FeignClient;
-import org.springframework.cloud.sleuth.instrument.reactor.TraceReactorAutoConfigurationAccessorConfiguration;
 import org.springframework.cloud.sleuth.instrument.web.TraceWebServletAutoConfiguration;
 import org.springframework.cloud.sleuth.util.ArrayListSpanReporter;
 import org.springframework.context.annotation.Bean;
@@ -163,11 +161,6 @@ public class WebClientTests {
 	@Autowired
 	MyRestTemplateCustomizer customizer;
 
-	@BeforeClass
-	public static void cleanup() {
-		TraceReactorAutoConfigurationAccessorConfiguration.close();
-	}
-
 	@After
 	public void close() {
 		this.reporter.clear();