feat(sdk-testing): Add W3CBaggagePropagator to test utils (#7056)

Author: edeweerd1A

sdk/testing/src/main/java/io/opentelemetry/sdk/testing/junit4/OpenTelemetryRule.java

@@ -7,8 +7,10 @@
 
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.baggage.propagation.W3CBaggagePropagator;
 import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
 import io.opentelemetry.context.propagation.ContextPropagators;
+import io.opentelemetry.context.propagation.TextMapPropagator;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import io.opentelemetry.sdk.logs.SdkLoggerProvider;
 import io.opentelemetry.sdk.logs.data.LogRecordData;
@@ -84,7 +86,11 @@ public static OpenTelemetryRule create() {
 
     OpenTelemetrySdk openTelemetry =
         OpenTelemetrySdk.builder()
-            .setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance()))
+            .setPropagators(
+                ContextPropagators.create(
+                    TextMapPropagator.composite(
+                        W3CTraceContextPropagator.getInstance(),
+                        W3CBaggagePropagator.getInstance())))
             .setTracerProvider(tracerProvider)
             .setMeterProvider(meterProvider)
             .setLoggerProvider(loggerProvider)

sdk/testing/src/main/java/io/opentelemetry/sdk/testing/junit5/OpenTelemetryExtension.java

@@ -9,8 +9,10 @@
 
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.baggage.propagation.W3CBaggagePropagator;
 import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
 import io.opentelemetry.context.propagation.ContextPropagators;
+import io.opentelemetry.context.propagation.TextMapPropagator;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import io.opentelemetry.sdk.logs.SdkLoggerProvider;
 import io.opentelemetry.sdk.logs.data.LogRecordData;
@@ -86,7 +88,11 @@ public static OpenTelemetryExtension create() {
 
     OpenTelemetrySdk openTelemetry =
         OpenTelemetrySdk.builder()
-            .setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance()))
+            .setPropagators(
+                ContextPropagators.create(
+                    TextMapPropagator.composite(
+                        W3CTraceContextPropagator.getInstance(),
+                        W3CBaggagePropagator.getInstance())))
             .setTracerProvider(tracerProvider)
             .setMeterProvider(meterProvider)
             .setLoggerProvider(loggerProvider)

sdk/testing/src/test/java/io/opentelemetry/sdk/testing/junit4/OpenTelemetryRuleTest.java

@@ -7,12 +7,20 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
+import io.opentelemetry.api.baggage.Baggage;
 import io.opentelemetry.api.common.Value;
 import io.opentelemetry.api.logs.Logger;
 import io.opentelemetry.api.metrics.LongCounter;
 import io.opentelemetry.api.metrics.Meter;
+import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.context.propagation.TextMapSetter;
 import io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -118,4 +126,31 @@ public void getLogRecordsAgain() {
         .satisfies(
             logRecordData -> assertThat(logRecordData.getBodyValue()).isEqualTo(Value.of("body")));
   }
+
+  @Test
+  public void baggageAndTracePropagation() {
+    OpenTelemetryRule rule = OpenTelemetryRule.create();
+    Span span = rule.getOpenTelemetry().getTracer("test").spanBuilder("test").startSpan();
+    try (Scope baggageScope = Baggage.builder().put("key", "value").build().makeCurrent();
+        Scope spanScope = span.makeCurrent()) {
+      Map<String, String> carrier = new HashMap<>();
+      rule.getOpenTelemetry()
+          .getPropagators()
+          .getTextMapPropagator()
+          .inject(Context.current(), carrier, new MapTextMapSetter());
+      assertThat(carrier).containsEntry("baggage", "key=value");
+      assertThat(carrier).containsKey("traceparent");
+    } finally {
+      span.end();
+    }
+  }
+
+  public static class MapTextMapSetter implements TextMapSetter<Map<String, String>> {
+    @Override
+    public void set(@Nullable Map<String, String> carrier, String key, String value) {
+      if (carrier != null) {
+        carrier.put(key, value);
+      }
+    }
+  }
 }

sdk/testing/src/test/java/io/opentelemetry/sdk/testing/junit5/OpenTelemetryExtensionTest.java

@@ -10,18 +10,24 @@
 
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.baggage.Baggage;
 import io.opentelemetry.api.common.Value;
 import io.opentelemetry.api.logs.Logger;
 import io.opentelemetry.api.metrics.LongCounter;
 import io.opentelemetry.api.metrics.Meter;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.context.propagation.TextMapSetter;
 import io.opentelemetry.sdk.trace.data.SpanData;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
+import javax.annotation.Nullable;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
@@ -231,4 +237,32 @@ void afterAll() {
     assertThat(extension.getMetrics()).isEmpty();
     assertThat(extension.getSpans()).isEmpty();
   }
+
+  @Test
+  void baggageAndTracePropagation() {
+    OpenTelemetryExtension extension = OpenTelemetryExtension.create();
+    Span span = extension.getOpenTelemetry().getTracer("test").spanBuilder("test").startSpan();
+    try (Scope baggageScope = Baggage.builder().put("key", "value").build().makeCurrent();
+        Scope spanScope = span.makeCurrent()) {
+      Map<String, String> carrier = new HashMap<>();
+      extension
+          .getOpenTelemetry()
+          .getPropagators()
+          .getTextMapPropagator()
+          .inject(Context.current(), carrier, new MapTextMapSetter());
+      assertThat(carrier).containsEntry("baggage", "key=value");
+      assertThat(carrier).containsKey("traceparent");
+    } finally {
+      span.end();
+    }
+  }
+
+  public static class MapTextMapSetter implements TextMapSetter<Map<String, String>> {
+    @Override
+    public void set(@Nullable Map<String, String> carrier, String key, String value) {
+      if (carrier != null) {
+        carrier.put(key, value);
+      }
+    }
+  }
 }