From e9d55a3af0b501f0fc994f5b134bf3ce0ac62872 Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <marcin.grzejszczak@gmail.com>
Date: Fri, 2 Jul 2021 13:13:42 +0200
Subject: [PATCH] Added awaitility to benchmarks

---
 benchmarks/pom.xml                            |  9 ++++-----
 .../SleuthBenchmarkingStreamApplication.java  | 20 +++++++++++--------
 .../jmh/stream/MicroBenchmarkStreamTests.java |  5 +++--
 3 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/benchmarks/pom.xml b/benchmarks/pom.xml
index 50c51b9ed..0972ea89a 100644
--- a/benchmarks/pom.xml
+++ b/benchmarks/pom.xml
@@ -41,7 +41,6 @@
 		<okhttp.version>4.9.0</okhttp.version>
 		<microbenchmark-runner.version>0.2.0.RELEASE</microbenchmark-runner.version>
 		<jmh.version>1.26</jmh.version>
-		<spring-cloud-stream.version>3.1.4-SNAPSHOT</spring-cloud-stream.version>
 	</properties>
 
 	<dependencyManagement>
@@ -55,8 +54,8 @@
 			</dependency>
 			<dependency>
 				<groupId>org.springframework.cloud</groupId>
-				<artifactId>spring-cloud-stream-dependencies</artifactId>
-				<version>${spring-cloud-stream.version}</version>
+				<artifactId>spring-cloud-sleuth</artifactId>
+				<version>${project.version}</version>
 				<type>pom</type>
 				<scope>import</scope>
 			</dependency>
@@ -65,7 +64,7 @@
 
 	<dependencies>
 		<dependency>
-			<groupId>${project.groupId}</groupId>
+			<groupId>org.springframework.cloud</groupId>
 			<artifactId>spring-cloud-starter-sleuth</artifactId>
 		</dependency>
 		<dependency>
@@ -150,7 +149,7 @@
 		<dependency>
 			<groupId>org.awaitility</groupId>
 			<artifactId>awaitility</artifactId>
-			<scope>test</scope>
+			<scope>compile</scope>
 		</dependency>
 	</dependencies>
 
diff --git a/benchmarks/src/main/java/org/springframework/cloud/sleuth/benchmarks/app/stream/SleuthBenchmarkingStreamApplication.java b/benchmarks/src/main/java/org/springframework/cloud/sleuth/benchmarks/app/stream/SleuthBenchmarkingStreamApplication.java
index 32f1fa718..6c176269d 100644
--- a/benchmarks/src/main/java/org/springframework/cloud/sleuth/benchmarks/app/stream/SleuthBenchmarkingStreamApplication.java
+++ b/benchmarks/src/main/java/org/springframework/cloud/sleuth/benchmarks/app/stream/SleuthBenchmarkingStreamApplication.java
@@ -23,6 +23,7 @@ import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 
+import org.awaitility.Awaitility;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Flux;
@@ -60,15 +61,16 @@ public class SleuthBenchmarkingStreamApplication {
 		// "DECORATE_ON_LAST");
 		// System.setProperty("spring.sleuth.reactor.instrumentation-type", "MANUAL");
 		System.setProperty("spring.sleuth.reactor.instrumentation-type", "DECORATE_QUEUES");
+		System.setProperty("spring.sleuth.integration.enabled", "true");
 		System.setProperty("spring.sleuth.function.type", "DECORATE_QUEUES");
 		ConfigurableApplicationContext context = SpringApplication.run(SleuthBenchmarkingStreamApplication.class, args);
-		for (int i = 0; i < 1; i++) {
-			InputDestination input = context.getBean(InputDestination.class);
-			input.send(MessageBuilder.withPayload("hello".getBytes())
-					.setHeader("b3", "4883117762eb9420-4883117762eb9420-1").build());
-			log.info("Retrieving the message for tests");
-			OutputDestination output = context.getBean(OutputDestination.class);
-			Message<byte[]> message = output.receive(200L);
+		InputDestination input = context.getBean(InputDestination.class);
+		input.send(MessageBuilder.withPayload("hello".getBytes())
+				.setHeader("b3", "4883117762eb9420-4883117762eb9420-1").build());
+		log.info("Retrieving the message for tests");
+		OutputDestination output = context.getBean(OutputDestination.class);
+		Awaitility.await().untilAsserted( () -> {
+			Message<byte[]> message = output.receive(1L);
 			log.info("Got the message from output");
 			assertThat(message).isNotNull();
 			log.info("Message is not null");
@@ -77,7 +79,9 @@ public class SleuthBenchmarkingStreamApplication {
 			String b3 = message.getHeaders().get("b3", String.class);
 			log.info("Checking the b3 header [" + b3 + "]");
 			assertThat(b3).startsWith("4883117762eb9420");
-		}
+		});
+		context.close();
+		System.exit(0);
 	}
 
 	@Bean
diff --git a/benchmarks/src/test/java/org/springframework/cloud/sleuth/benchmarks/jmh/stream/MicroBenchmarkStreamTests.java b/benchmarks/src/test/java/org/springframework/cloud/sleuth/benchmarks/jmh/stream/MicroBenchmarkStreamTests.java
index 6ce4905df..be5393769 100644
--- a/benchmarks/src/test/java/org/springframework/cloud/sleuth/benchmarks/jmh/stream/MicroBenchmarkStreamTests.java
+++ b/benchmarks/src/test/java/org/springframework/cloud/sleuth/benchmarks/jmh/stream/MicroBenchmarkStreamTests.java
@@ -23,6 +23,7 @@ import java.util.concurrent.TimeUnit;
 
 import brave.Tracing;
 import jmh.mbr.junit5.Microbenchmark;
+import org.awaitility.Awaitility;
 import org.junit.platform.commons.annotation.Testable;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
@@ -114,12 +115,12 @@ public class MicroBenchmarkStreamTests {
 
 		void run() {
 			sendInputMessage();
-			assertThatOutputMessageGotReceived();
+			Awaitility.await().untilAsserted(this::assertThatOutputMessageGotReceived);
 		}
 
 		private void assertThatOutputMessageGotReceived() {
 			// System.out.println("Retrieving the message for tests");
-			Message<byte[]> message = output.receive(200L);
+			Message<byte[]> message = output.receive(1L);
 			// System.out.println("Got the message from output");
 			assertThat(message).isNotNull();
 			// System.out.println("Message is not null");