Polish BaggagePropagationIntegrationTests

Author: mhalbritter

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/tracing/BaggagePropagationIntegrationTests.java

@@ -23,6 +23,7 @@
 import io.micrometer.tracing.Span;
 import io.micrometer.tracing.Tracer;
 import io.opentelemetry.context.Context;
+import org.assertj.core.api.ThrowingConsumer;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -41,11 +42,13 @@
  * formats.
  *
  * @author Marcin Grzejszczak
+ * @author Moritz Halbritter
  */
 class BaggagePropagationIntegrationTests {
 
-	static final String COUNTRY_CODE = "country-code";
-	static final String BUSINESS_PROCESS = "bp";
+	private static final String COUNTRY_CODE = "country-code";
+
+	private static final String BUSINESS_PROCESS = "bp";
 
 	@BeforeEach
 	@AfterEach
@@ -61,7 +64,7 @@ void shouldSetEntriesToMdcFromSpanWithBaggage(AutoConfig autoConfig) {
 			Span span = createSpan(tracer);
 			assertThatTracingContextIsInitialized(autoConfig);
 			try (Tracer.SpanInScope scope = tracer.withSpan(span.start())) {
-				BaggageManager baggageManager = context.getBean(BaggageManager.class);
+				BaggageManager baggageManager = baggageManager(context);
 				try (BaggageInScope fo = baggageManager.createBaggageInScope(span.context(), COUNTRY_CODE, "FO");
 						BaggageInScope alm = baggageManager.createBaggageInScope(span.context(), BUSINESS_PROCESS,
 								"ALM")) {
@@ -73,8 +76,7 @@ void shouldSetEntriesToMdcFromSpanWithBaggage(AutoConfig autoConfig) {
 			finally {
 				span.end();
 			}
-
-			assertThatMdcContainsUnsetTraceId();
+			assertThatMdcContainsUnsetTraceId(autoConfig);
 			assertThat(MDC.get(COUNTRY_CODE)).isNull();
 			assertThat(MDC.get(BUSINESS_PROCESS)).isNull();
 		});
@@ -88,25 +90,22 @@ void shouldRemoveEntriesFromMdcForNullSpan(AutoConfig autoConfig) {
 			Span span = createSpan(tracer);
 			assertThatTracingContextIsInitialized(autoConfig);
 			try (Tracer.SpanInScope scope = tracer.withSpan(span.start())) {
-				try (BaggageInScope fo = context.getBean(BaggageManager.class)
-					.createBaggageInScope(span.context(), COUNTRY_CODE, "FO")) {
-
+				try (BaggageInScope fo = baggageManager(context).createBaggageInScope(span.context(), COUNTRY_CODE,
+						"FO")) {
 					assertThat(MDC.get("traceId")).isEqualTo(span.context().traceId());
 					assertThat(MDC.get(COUNTRY_CODE)).isEqualTo("FO");
-
 					try (Tracer.SpanInScope scope2 = tracer.withSpan(null)) {
-						assertThatMdcContainsUnsetTraceId();
+						assertThatMdcContainsUnsetTraceId(autoConfig);
 						assertThat(MDC.get(COUNTRY_CODE)).isNull();
 					}
-
 					assertThat(MDC.get("traceId")).isEqualTo(span.context().traceId());
 					assertThat(MDC.get(COUNTRY_CODE)).isEqualTo("FO");
 				}
 			}
 			finally {
 				span.end();
 			}
-			assertThatMdcContainsUnsetTraceId();
+			assertThatMdcContainsUnsetTraceId(autoConfig);
 			assertThat(MDC.get(COUNTRY_CODE)).isNull();
 		});
 	}
@@ -119,22 +118,28 @@ private Tracer tracer(ApplicationContext context) {
 		return context.getBean(Tracer.class);
 	}
 
+	private BaggageManager baggageManager(ApplicationContext context) {
+		return context.getBean(BaggageManager.class);
+	}
+
 	private void assertThatTracingContextIsInitialized(AutoConfig autoConfig) {
-		if (autoConfig == AutoConfig.OTEL_B3) {
+		if (autoConfig.isOtel()) {
 			assertThat(Context.current()).isEqualTo(Context.root());
 		}
 	}
 
-	private void assertThatMdcContainsUnsetTraceId() {
-		assertThat(isInvalidBraveTraceId() || isInvalidOtelTraceId()).isTrue();
-	}
-
-	private boolean isInvalidBraveTraceId() {
-		return MDC.get("traceId") == null;
-	}
-
-	private boolean isInvalidOtelTraceId() {
-		return MDC.get("traceId").equals("00000000000000000000000000000000");
+	private void assertThatMdcContainsUnsetTraceId(AutoConfig autoConfig) {
+		boolean eitherOtelOrBrave = autoConfig.isOtel() || autoConfig.isBrave();
+		assertThat(eitherOtelOrBrave).isTrue();
+		if (autoConfig.isOtel()) {
+			ThrowingConsumer<String> isNull = (traceId) -> assertThat(traceId).isNull();
+			ThrowingConsumer<String> isZero = (traceId) -> assertThat(traceId)
+				.isEqualTo("00000000000000000000000000000000");
+			assertThat(MDC.get("traceId")).satisfiesAnyOf(isNull, isZero);
+		}
+		if (autoConfig.isBrave()) {
+			assertThat(MDC.get("traceId")).isNull();
+		}
 	}
 
 	enum AutoConfig implements Supplier<ApplicationContextRunner> {
@@ -227,6 +232,14 @@ public ApplicationContextRunner get() {
 							"management.tracing.baggage.remote-fields=x-vcap-request-id,country-code,bp",
 							"management.tracing.baggage.correlation.fields=country-code,bp");
 			}
+		};
+
+		boolean isOtel() {
+			return name().startsWith("OTEL_");
+		}
+
+		boolean isBrave() {
+			return name().startsWith("BRAVE_");
 		}
 
 	}