adding more tests

Author: rachelyangdog

dd-java-agent/instrumentation/okhttp-3/src/test/groovy/OkHttp3Test.groovy

@@ -78,6 +78,65 @@ abstract class OkHttp3Test extends HttpClientTest {
     method = "GET"
     url = server.address.resolve(path)
   }
+
+  def "baggage span tags are properly added"() {
+    setup:
+    // W3C Baggage header format: key1=value1,key2=value2,key3=value3
+    def baggageHeader = "user.id=bark,session.id=test-sess1,account.id=fetch,language=en"
+    def sentHeaders = [:]
+
+    when:
+    def status
+    // Capture the headers that OkHttp3 sends
+    def interceptor = new Interceptor() {
+        @Override
+        Response intercept(Chain chain) throws IOException {
+          def request = chain.request()
+          // Capture all headers sent
+          request.headers().names().each { headerName ->
+            sentHeaders[headerName] = request.header(headerName)
+          }
+          return chain.proceed(request)
+        }
+      }
+
+    def testClient = new OkHttpClient.Builder()
+      .connectTimeout(CONNECT_TIMEOUT_MS, TimeUnit.MILLISECONDS)
+      .readTimeout(READ_TIMEOUT_MS, TimeUnit.MILLISECONDS)
+      .writeTimeout(READ_TIMEOUT_MS, TimeUnit.MILLISECONDS)
+      .addInterceptor(interceptor)
+      .build()
+
+    def request = new Request.Builder()
+      .url(server.address.resolve("/success").toURL())
+      .header("baggage", baggageHeader)  // Pass baggage directly in header
+      .get()
+      .build()
+    def response = testClient.newCall(request).execute()
+    status = response.code()
+
+    then:
+    status == 200
+    // Verify baggage header was sent
+    sentHeaders["baggage"] == baggageHeader
+    sentHeaders["baggage"].contains("user.id=bark")
+    sentHeaders["baggage"].contains("session.id=test-sess1")
+    sentHeaders["baggage"].contains("account.id=fetch")
+    sentHeaders["baggage"].contains("language=en")
+
+    // Verify the resulting span has the correct baggage tags (only default configured keys)
+    assertTraces(1) {
+      trace(1) {
+        clientSpan(it, null, "GET", false, false, server.address.resolve("/success"), 200, false, null, false, [
+          // Should have baggage tags for keys in default configuration
+          "baggage.user.id": "bark",
+          "baggage.session.id": "test-sess1",
+          "baggage.account.id": "fetch",
+          // "baggage.language" should NOT be present since it's not in default config
+        ])
+      }
+    }
+  }
 }
 
 @Timeout(5)

dd-trace-core/src/test/groovy/datadog/trace/core/CoreSpanBuilderTest.groovy

@@ -4,11 +4,13 @@ import static datadog.trace.api.DDTags.DJM_ENABLED
 import static datadog.trace.api.DDTags.DSM_ENABLED
 import static datadog.trace.api.DDTags.PROFILING_ENABLED
 import static datadog.trace.api.DDTags.SCHEMA_VERSION_TAG_KEY
+import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_TAG_KEYS
 
 import datadog.trace.api.Config
 import datadog.trace.api.DDSpanId
 import datadog.trace.api.DDTraceId
 import datadog.trace.api.gateway.RequestContextSlot
+import datadog.trace.api.sampling.PrioritySampling
 import datadog.trace.api.naming.SpanNaming
 import datadog.trace.api.sampling.PrioritySampling
 import datadog.trace.bootstrap.instrumentation.api.AgentScope
@@ -491,6 +493,31 @@ class CoreSpanBuilderTest extends DDCoreSpecification {
     span1.finish()
   }
 
+  def "buildSpan should add baggage tags with different configurations"() {
+    setup:
+    injectSysConfig(TRACE_BAGGAGE_TAG_KEYS, baggageTagKeysConfig)
+    def baggage = ["user.id": "alice", "session.id": "123", "region": "us-west-1", "env": "production"]
+    def tagContext = new TagContext(null, null, null, baggage, PrioritySampling.UNSET, null, DATADOG, DDTraceId.ZERO)
+
+    when:
+    def span = tracer.buildSpan("test", "test-op")
+      .asChildOf(tagContext)
+      .start()
+
+    then:
+    // Filter span tags to only check baggage tags (those starting with "baggage.")
+    def actualBaggageTags = span.tags.findAll { key, value -> key.startsWith("baggage.") }
+    actualBaggageTags == expectedBaggageTags
+
+    cleanup:
+    span.finish()
+
+    where:
+    baggageTagKeysConfig | expectedBaggageTags
+    "user.id"            | ["baggage.user.id": "alice"]
+    "user.id,session.id" | ["baggage.user.id": "alice", "baggage.session.id": "123"]
+  }
+
   def productTags() {
     def productTags = [
       (PROFILING_ENABLED) : Config.get().isProfilingEnabled() ? 1 : 0

dd-trace-core/src/test/groovy/datadog/trace/core/CoreTracerTest.groovy

@@ -38,6 +38,7 @@ import static datadog.trace.api.config.TracerConfig.HEADER_TAGS
 import static datadog.trace.api.config.TracerConfig.PRIORITY_SAMPLING
 import static datadog.trace.api.config.TracerConfig.SERVICE_MAPPING
 import static datadog.trace.api.config.TracerConfig.SPAN_TAGS
+import static datadog.trace.api.config.TracerConfig.TRACE_BAGGAGE_TAG_KEYS
 import static datadog.trace.api.config.TracerConfig.WRITER_TYPE
 
 @Timeout(10)
@@ -640,11 +641,31 @@ class CoreTracerTest extends DDCoreSpecification {
     "service" | "env"  | "service_2"   | "env_2"
   }
 
-  def "test mapBaggageTags"() {
+  def "test mapBaggageTags default"() {
     when:
-    def tags = CoreTracer.mapBaggageTags(["user.id": "doggo", "foo": "bar"])
+    def tags = CoreTracer.mapBaggageTags(["user.id": "doggo", "account.id": "test", "session.id":"1234", "region":"us-east-1"])
     then:
-    tags == ["baggage.user.id": "doggo"]
+    tags == ["baggage.user.id": "doggo", "baggage.account.id": "test", "baggage.session.id": "1234"]
+  }
+
+  def "test mapBaggageTags with different configurations"() {
+    setup:
+    injectSysConfig(TRACE_BAGGAGE_TAG_KEYS, baggageTagKeysConfig)
+
+    when:
+    def tags = CoreTracer.mapBaggageTags(["user.id": "doggo", "foo": "bar", "language":"en"])
+
+    then:
+    tags == expectedTags
+
+    where:
+    baggageTagKeysConfig | expectedTags
+    "*"                  | ["baggage.user.id": "doggo", "baggage.foo": "bar", "baggage.language": "en"]
+    " "                  | [:]
+    "system.id"          | [:]
+    "user.id"            | ["baggage.user.id": "doggo"]
+    "foo"                | ["baggage.foo": "bar"]
+    "user.id,foo"        | ["baggage.user.id": "doggo", "baggage.foo": "bar"]
   }
 }