Handle nulls in noop / multi (#2939)

* Handle nulls in noop / multi * Remove old file
open-telemetry · May 7, 2021 · 0b8131c · 0b8131c
1 parent bd6cfd8
commit 0b8131c
Show file tree

Hide file tree

Showing 4 changed files with 81 additions and 29 deletions.
diff --git a/context/src/main/java/io/opentelemetry/context/propagation/MultiTextMapPropagator.java b/context/src/main/java/io/opentelemetry/context/propagation/MultiTextMapPropagator.java
@@ -45,13 +45,22 @@ private static List<String> getAllFields(TextMapPropagator[] textPropagators) {
 
   @Override
   public <C> void inject(Context context, @Nullable C carrier, TextMapSetter<C> setter) {
+    if (context == null || setter == null) {
+      return;
+    }
     for (TextMapPropagator textPropagator : textPropagators) {
       textPropagator.inject(context, carrier, setter);
     }
   }
 
   @Override
   public <C> Context extract(Context context, @Nullable C carrier, TextMapGetter<C> getter) {
+    if (context == null) {
+      return Context.root();
+    }
+    if (getter == null) {
+      return context;
+    }
     for (TextMapPropagator textPropagator : textPropagators) {
       context = textPropagator.extract(context, carrier, getter);
     }

diff --git a/context/src/main/java/io/opentelemetry/context/propagation/NoopTextMapPropagator.java b/context/src/main/java/io/opentelemetry/context/propagation/NoopTextMapPropagator.java
@@ -27,6 +27,9 @@ public <C> void inject(Context context, @Nullable C carrier, TextMapSetter<C> se
 
   @Override
   public <C> Context extract(Context context, @Nullable C carrier, TextMapGetter<C> getter) {
+    if (context == null) {
+      return Context.root();
+    }
     return context;
   }
 }
diff --git a/context/src/test/java/io/opentelemetry/context/propagation/MultiTextMapPropagatorTest.java b/context/src/test/java/io/opentelemetry/context/propagation/MultiTextMapPropagatorTest.java
@@ -12,9 +12,12 @@
 import static org.mockito.Mockito.when;
 
 import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import javax.annotation.Nullable;
@@ -26,6 +29,8 @@
 @ExtendWith(MockitoExtension.class)
 class MultiTextMapPropagatorTest {
 
+  private static final ContextKey<String> KEY = ContextKey.named("key");
+
   @Mock private TextMapPropagator propagator1;
   @Mock private TextMapPropagator propagator2;
   @Mock private TextMapPropagator propagator3;
@@ -93,19 +98,6 @@ void inject_allDelegated() {
     verify(propagator3).inject(context, carrier, setter);
   }
 
-  @Test
-  void inject_nullContextAllDelegated() {
-    Map<String, String> carrier = new HashMap<>();
-    Context context = null;
-    TextMapSetter<Map<String, String>> setter = Map::put;
-
-    TextMapPropagator prop = new MultiTextMapPropagator(propagator1, propagator2, propagator3);
-    prop.inject(context, carrier, setter);
-    verify(propagator1).inject(context, carrier, setter);
-    verify(propagator2).inject(context, carrier, setter);
-    verify(propagator3).inject(context, carrier, setter);
-  }
-
   @Test
   void extract_noPropagators() {
     Map<String, String> carrier = new HashMap<>();
@@ -147,14 +139,33 @@ void extract_notFound() {
 
   @Test
   void extract_nullContext() {
-    Map<String, String> carrier = new HashMap<>();
-    Context context = null;
-    when(propagator1.extract(context, carrier, getter)).thenReturn(context);
-    when(propagator2.extract(context, carrier, getter)).thenReturn(context);
+    assertThat(
+            new MultiTextMapPropagator(propagator1, propagator2)
+                .extract(null, Collections.emptyMap(), getter))
+        .isSameAs(Context.root());
+  }
 
-    TextMapPropagator prop = new MultiTextMapPropagator(propagator1, propagator2);
-    Context result = prop.extract(context, carrier, getter);
+  @Test
+  void extract_nullGetter() {
+    Context context = Context.current().with(KEY, "treasure");
+    assertThat(
+            new MultiTextMapPropagator(propagator1, propagator2)
+                .extract(context, Collections.emptyMap(), null))
+        .isSameAs(context);
+  }
 
-    assertThat(result).isSameAs(context);
+  @Test
+  void inject_nullContext() {
+    Map<String, String> carrier = new LinkedHashMap<>();
+    new MultiTextMapPropagator(propagator1, propagator2).inject(null, carrier, Map::put);
+    assertThat(carrier).isEmpty();
+  }
+
+  @Test
+  void inject_nullSetter() {
+    Map<String, String> carrier = new LinkedHashMap<>();
+    Context context = Context.current().with(KEY, "treasure");
+    new MultiTextMapPropagator(propagator1, propagator2).inject(context, carrier, null);
+    assertThat(carrier).isEmpty();
   }
 }
diff --git a/context/src/test/java/io/opentelemetry/context/propagation/NoopTextMapPropagatorTest.java b/context/src/test/java/io/opentelemetry/context/propagation/NoopTextMapPropagatorTest.java
@@ -8,12 +8,18 @@
 import static org.assertj.core.api.Assertions.assertThat;
 
 import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.Map;
 import javax.annotation.Nullable;
 import org.junit.jupiter.api.Test;
 
 class NoopTextMapPropagatorTest {
 
+  private static final ContextKey<String> KEY = ContextKey.named("key");
+
   @Test
   void noopFields() {
     assertThat(TextMapPropagator.noop().fields()).isEmpty();
@@ -23,28 +29,51 @@ void noopFields() {
   void extract_contextUnchanged() {
     Context input = Context.current();
     Context result =
-        TextMapPropagator.noop().extract(input, new HashMap<>(), new HashMapTextMapGetter());
+        TextMapPropagator.noop().extract(input, new HashMap<>(), MapTextMapGetter.INSTANCE);
     assertThat(result).isSameAs(input);
   }
 
   @Test
   void extract_nullContext() {
-    Context input = null;
-    Context result =
-        TextMapPropagator.noop().extract(input, new HashMap<>(), new HashMapTextMapGetter());
-    assertThat(result).isSameAs(input);
+    assertThat(
+            TextMapPropagator.noop()
+                .extract(null, Collections.emptyMap(), MapTextMapGetter.INSTANCE))
+        .isSameAs(Context.root());
   }
 
-  private static class HashMapTextMapGetter
-      implements TextMapGetter<HashMap<? extends Object, ? extends Object>> {
+  @Test
+  void extract_nullGetter() {
+    Context context = Context.current().with(KEY, "treasure");
+    assertThat(TextMapPropagator.noop().extract(context, Collections.emptyMap(), null))
+        .isSameAs(context);
+  }
+
+  @Test
+  void inject_nullContext() {
+    Map<String, String> carrier = new LinkedHashMap<>();
+    TextMapPropagator.noop().inject(null, carrier, Map::put);
+    assertThat(carrier).isEmpty();
+  }
+
+  @Test
+  void inject_nullSetter() {
+    Map<String, String> carrier = new LinkedHashMap<>();
+    Context context = Context.current().with(KEY, "treasure");
+    TextMapPropagator.noop().inject(context, carrier, null);
+    assertThat(carrier).isEmpty();
+  }
+
+  enum MapTextMapGetter implements TextMapGetter<Map<? extends Object, ? extends Object>> {
+    INSTANCE;
+
     @Override
-    public Iterable<String> keys(HashMap<?, ?> carrier) {
+    public Iterable<String> keys(Map<?, ?> carrier) {
       return null;
     }
 
     @Nullable
     @Override
-    public String get(@Nullable HashMap<?, ?> carrier, String key) {
+    public String get(@Nullable Map<?, ?> carrier, String key) {
       return null;
     }
   }