writing unit tests

Author: mhlidd

dd-java-agent/instrumentation/armeria/armeria-grpc-0.84/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/GrpcInjectAdapter.java

@@ -12,7 +12,7 @@ public final class GrpcInjectAdapter implements CarrierSetter<Metadata> {
   @Override
   public void set(final Metadata carrier, final String key, final String value) {
     Metadata.Key<String> metadataKey = Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER);
-    if(!key.startsWith("ot-baggage-*") || !carrier.containsKey(metadataKey)){
+    if (!key.startsWith("ot-baggage-") || !carrier.containsKey(metadataKey)) {
       carrier.put(metadataKey, value);
     }
   }

dd-java-agent/instrumentation/armeria/armeria-grpc-0.84/src/test/groovy/GrpcInjectAdapterTest.groovy

@@ -0,0 +1,47 @@
+import datadog.trace.agent.test.naming.VersionedNamingTestBase
+import io.grpc.Metadata
+import static datadog.trace.instrumentation.grpc.client.GrpcInjectAdapter.SETTER
+
+class GrpcInjectAdapterTest extends VersionedNamingTestBase {
+  def "carrier set is called only once per unique ot-baggage-* key"() {
+    setup:
+    def carrier = new Metadata()
+
+    def baggage = [
+      ["ot-baggage-foo", "v1"],
+      ["ot-baggage-foo", "v2"],
+      ["ot-baggage-bar", "v3"]
+    ]
+
+    when:
+    baggage.each { pair ->
+      def (key, value) = pair
+      SETTER.set(carrier, key, value)
+    }
+
+    then:
+    carrier.headerCount() == 2
+    carrier.get(getKey("ot-baggage-foo")) == "v1" // first value wins
+    carrier.get(getKey("ot-baggage-bar")) == "v3"
+  }
+
+  Metadata.Key<String> getKey(String key){
+    Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER)
+  }
+
+
+  @Override
+  int version() {
+    return 0
+  }
+
+  @Override
+  String service() {
+    return null
+  }
+
+  @Override
+  String operation() {
+    return null
+  }
+}

dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/client/GrpcInjectAdapter.java

@@ -10,6 +10,9 @@ public final class GrpcInjectAdapter implements CarrierSetter<Metadata> {
 
   @Override
   public void set(final Metadata carrier, final String key, final String value) {
-    carrier.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value);
+    Metadata.Key<String> metadataKey = Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER);
+    if (!key.startsWith("ot-baggage-") || !carrier.containsKey(metadataKey)) {
+      carrier.put(metadataKey, value);
+    }
   }
 }

dd-java-agent/instrumentation/grpc-1.5/src/test/groovy/GrpcInjectAdapterTest.groovy

@@ -0,0 +1,47 @@
+import datadog.trace.agent.test.naming.VersionedNamingTestBase
+import io.grpc.Metadata
+import static datadog.trace.instrumentation.grpc.client.GrpcInjectAdapter.SETTER
+
+class GrpcInjectAdapterTest extends VersionedNamingTestBase {
+  def "carrier set is called only once per unique ot-baggage-* key"() {
+    setup:
+    def carrier = new Metadata()
+
+    def baggage = [
+      ["ot-baggage-foo", "v1"],
+      ["ot-baggage-foo", "v2"],
+      ["ot-baggage-bar", "v3"]
+    ]
+
+    when:
+    baggage.each { pair ->
+      def (key, value) = pair
+      SETTER.set(carrier, key, value)
+    }
+
+    then:
+    carrier.headerCount() == 2
+    carrier.get(getKey("ot-baggage-foo")) == "v1" // first value wins
+    carrier.get(getKey("ot-baggage-bar")) == "v3"
+  }
+
+  Metadata.Key<String> getKey(String key){
+    Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER)
+  }
+
+
+  @Override
+  int version() {
+    return 0
+  }
+
+  @Override
+  String service() {
+    return null
+  }
+
+  @Override
+  String operation() {
+    return null
+  }
+}

dd-trace-core/src/test/groovy/datadog/trace/core/propagation/HttpInjectorTest.groovy

@@ -218,35 +218,6 @@ class HttpInjectorTest extends DDCoreSpecification {
     style << [DATADOG, TRACECONTEXT, HAYSTACK]
   }
 
-  def "inject ot-baggage-* in http headers only once"() {
-    setup:
-    Config config = Mock(Config)
-    def baggage = [
-      'k1': 'v1',
-      'k2': 'v2',
-    ]
-    def mapping = baggage.keySet().collectEntries { [(it):it]} as Map<String, String>
-    def injector = HttpCodec.createInjector(config, [DATADOG, TRACECONTEXT].toSet(), mapping)
-    def traceId = DDTraceId.ONE
-    def spanId = 2
-    def writer = new ListWriter()
-    def tracer = tracerBuilder().writer(writer).build()
-    final DDSpanContext mockedContext = mockedContext(tracer, traceId, spanId, UNSET, null, baggage)
-    final Map<String, String> carrier = Mock()
-    mockedContext.beginEndToEnd()
-
-    when:
-    injector.inject(mockedContext, carrier, MapSetter.INSTANCE)
-
-    then:
-    1 * carrier.put('k1', 'v1')
-    1 * carrier.put('k2', 'v2')
-    1 * carrier.put('ot-baggage-t0', "${(long) (mockedContext.endToEndStartTime / 1000000L)}")
-
-    cleanup:
-    tracer.close()
-  }
-
   static DDSpanContext mockedContext(CoreTracer tracer, DDTraceId traceId, long spanId, int samplingPriority, String origin, Map<String, String> baggage) {
     return new DDSpanContext(
       traceId,