[test] More tests for Reactor Core feature Hooks.enableAutomaticConte…

…xtPropagation() (#2926)
reactor · Oct 5, 2023 · bb8a50a · bb8a50a
1 parent ac7f7ba
commit bb8a50a
Showing 1 changed file with 70 additions and 18 deletions.
diff --git a/reactor-netty-http/src/contextPropagationTest/java/reactor/netty/ContextPropagationTest.java b/reactor-netty-http/src/contextPropagationTest/java/reactor/netty/ContextPropagationTest.java
@@ -144,24 +144,17 @@ void testAutomaticContextPropagation(HttpClient client) {
 
 			TestThreadLocalHolder.value("First");
 
-			AtomicReference<String> threadLocal = new AtomicReference<>();
-			client.port(disposableServer.port())
-			      .wiretap(true)
-			      .post()
-			      .uri("/")
-			      .send(ByteBufMono.fromString(Mono.just("test")))
-			      .responseContent()
-			      .aggregate()
-			      .asString()
-			      .doOnNext(s -> threadLocal.set(TestThreadLocalHolder.value()))
-			      .block(Duration.ofSeconds(5));
-
-			if (enableAutomaticContextPropagation) {
-				assertThat(threadLocal.get()).isNotNull().isEqualTo("First");
-			}
-			else {
-				assertThat(threadLocal.get()).isNull();
-			}
+			HttpClient.ResponseReceiver<?> responseReceiver =
+					client.port(disposableServer.port())
+					      .wiretap(true)
+					      .post()
+					      .uri("/")
+					      .send(ByteBufMono.fromString(Mono.just("test")));
+
+			response(responseReceiver, enableAutomaticContextPropagation);
+			responseConnection(responseReceiver, enableAutomaticContextPropagation);
+			responseContent(responseReceiver, enableAutomaticContextPropagation);
+			responseSingle(responseReceiver, enableAutomaticContextPropagation);
 		}
 		finally {
 			TestThreadLocalHolder.reset();
@@ -213,6 +206,65 @@ static Object[] httpClientCombinations() {
 		};
 	}
 
+	static void response(HttpClient.ResponseReceiver<?> responseReceiver, boolean enableAutomaticContextPropagation) {
+		AtomicReference<String> threadLocal = new AtomicReference<>();
+		responseReceiver.response()
+		                .doOnNext(s -> threadLocal.set(TestThreadLocalHolder.value()))
+		                .block(Duration.ofSeconds(5));
+
+		if (enableAutomaticContextPropagation) {
+			assertThat(threadLocal.get()).isNotNull().isEqualTo("First");
+		}
+		else {
+			assertThat(threadLocal.get()).isNull();
+		}
+	}
+
+	static void responseConnection(HttpClient.ResponseReceiver<?> responseReceiver, boolean enableAutomaticContextPropagation) {
+		AtomicReference<String> threadLocal = new AtomicReference<>();
+		responseReceiver.responseConnection((res, conn) -> conn.inbound().receive().aggregate().asString())
+		                .next()
+		                .doOnNext(s -> threadLocal.set(TestThreadLocalHolder.value()))
+		                .block(Duration.ofSeconds(5));
+
+		if (enableAutomaticContextPropagation) {
+			assertThat(threadLocal.get()).isNotNull().isEqualTo("First");
+		}
+		else {
+			assertThat(threadLocal.get()).isNull();
+		}
+	}
+
+	static void responseContent(HttpClient.ResponseReceiver<?> responseReceiver, boolean enableAutomaticContextPropagation) {
+		AtomicReference<String> threadLocal = new AtomicReference<>();
+		responseReceiver.responseContent()
+		                .aggregate()
+		                .asString()
+		                .doOnNext(s -> threadLocal.set(TestThreadLocalHolder.value()))
+		                .block(Duration.ofSeconds(5));
+
+		if (enableAutomaticContextPropagation) {
+			assertThat(threadLocal.get()).isNotNull().isEqualTo("First");
+		}
+		else {
+			assertThat(threadLocal.get()).isNull();
+		}
+	}
+
+	static void responseSingle(HttpClient.ResponseReceiver<?> responseReceiver, boolean enableAutomaticContextPropagation) {
+		AtomicReference<String> threadLocal = new AtomicReference<>();
+		responseReceiver.responseSingle((res, bytes) -> bytes.asString())
+		                .doOnNext(s -> threadLocal.set(TestThreadLocalHolder.value()))
+		                .block(Duration.ofSeconds(5));
+
+		if (enableAutomaticContextPropagation) {
+			assertThat(threadLocal.get()).isNotNull().isEqualTo("First");
+		}
+		else {
+			assertThat(threadLocal.get()).isNull();
+		}
+	}
+
 	static void sendRequest(HttpClient client, String expectation) {
 		client.post()
 		      .uri("/")