Improve tests for OpenCensusTraceContextDataInjector. (#1422)

This commit adds unit tests for
OpenCensusTraceContextDataInjector.rawContextData().  It also makes several
other minor improvements to the Log4j log correlation tests:

- Puts TestSpan in a separate file for reuse.
- Explicitly sets the SpanSelection in tests, where possible.
- Simplifies the tests by calling ReadOnlyStringMap.toMap().

4 files changed, 125 insertions, 49 deletions

diff --git a/contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjector.java b/contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjector.java
index c00a11a2..38b18826 100644
--- a/contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjector.java
+++ b/contrib/log_correlation/log4j2/src/main/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjector.java
@@ -143,7 +143,8 @@ public final class OpenCensusTraceContextDataInjector implements ContextDataInje
     this(lookUpSpanSelectionProperty());
   }
 
-  private OpenCensusTraceContextDataInjector(SpanSelection spanSelection) {
+  // visible for testing
+  OpenCensusTraceContextDataInjector(SpanSelection spanSelection) {
     this.spanSelection = spanSelection;
   }
 
diff --git a/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/AbstractOpenCensusLog4jLogCorrelationTest.java b/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/AbstractOpenCensusLog4jLogCorrelationTest.java
index b3811bc7..93ad85e0 100644
--- a/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/AbstractOpenCensusLog4jLogCorrelationTest.java
+++ b/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/AbstractOpenCensusLog4jLogCorrelationTest.java
@@ -19,18 +19,11 @@ package io.opencensus.contrib.logcorrelation.log4j2;
 import io.opencensus.common.Function;
 import io.opencensus.common.Scope;
 import io.opencensus.contrib.logcorrelation.log4j2.OpenCensusTraceContextDataInjector.SpanSelection;
-import io.opencensus.trace.Annotation;
-import io.opencensus.trace.AttributeValue;
-import io.opencensus.trace.EndSpanOptions;
-import io.opencensus.trace.Link;
-import io.opencensus.trace.Span;
 import io.opencensus.trace.SpanContext;
 import io.opencensus.trace.Tracer;
 import io.opencensus.trace.Tracestate;
 import io.opencensus.trace.Tracing;
 import java.io.StringWriter;
-import java.util.EnumSet;
-import java.util.Map;
 import org.apache.logging.log4j.Level;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.core.Appender;
@@ -101,22 +94,4 @@ abstract class AbstractOpenCensusLog4jLogCorrelationTest {
       scope.close();
     }
   }
-
-  private static final class TestSpan extends Span {
-    TestSpan(SpanContext context) {
-      super(context, EnumSet.of(Options.RECORD_EVENTS));
-    }
-
-    @Override
-    public void end(EndSpanOptions options) {}
-
-    @Override
-    public void addLink(Link link) {}
-
-    @Override
-    public void addAnnotation(Annotation annotation) {}
-
-    @Override
-    public void addAnnotation(String description, Map<String, AttributeValue> attributes) {}
-  }
 }
diff --git a/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjectorTest.java b/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjectorTest.java
index e027dae3..3b704058 100644
--- a/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjectorTest.java
+++ b/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/OpenCensusTraceContextDataInjectorTest.java
@@ -18,13 +18,19 @@ package io.opencensus.contrib.logcorrelation.log4j2;
 
 import static com.google.common.truth.Truth.assertThat;
 
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
+import io.opencensus.common.Scope;
 import io.opencensus.contrib.logcorrelation.log4j2.OpenCensusTraceContextDataInjector.SpanSelection;
+import io.opencensus.trace.SpanContext;
+import io.opencensus.trace.SpanId;
+import io.opencensus.trace.TraceId;
+import io.opencensus.trace.TraceOptions;
+import io.opencensus.trace.Tracer;
+import io.opencensus.trace.Tracestate;
+import io.opencensus.trace.Tracing;
 import java.util.Collections;
-import java.util.Map;
+import org.apache.logging.log4j.ThreadContext;
 import org.apache.logging.log4j.core.config.Property;
-import org.apache.logging.log4j.util.BiConsumer;
 import org.apache.logging.log4j.util.SortedArrayStringMap;
 import org.apache.logging.log4j.util.StringMap;
 import org.junit.Test;
@@ -34,6 +40,9 @@ import org.junit.runners.JUnit4;
 /** Tests for {@link OpenCensusTraceContextDataInjector}. */
 @RunWith(JUnit4.class)
 public final class OpenCensusTraceContextDataInjectorTest {
+  static final Tracestate EMPTY_TRACESTATE = Tracestate.builder().build();
+
+  private final Tracer tracer = Tracing.getTracer();
 
   @Test
   @SuppressWarnings("TruthConstantAsserts")
@@ -94,13 +103,13 @@ public final class OpenCensusTraceContextDataInjectorTest {
   @Test
   public void insertConfigurationProperties() {
     assertThat(
-            toMap(
-                new OpenCensusTraceContextDataInjector()
-                    .injectContextData(
-                        Lists.newArrayList(
-                            Property.createProperty("property1", "value1"),
-                            Property.createProperty("property2", "value2")),
-                        new SortedArrayStringMap())))
+            new OpenCensusTraceContextDataInjector(SpanSelection.ALL_SPANS)
+                .injectContextData(
+                    Lists.newArrayList(
+                        Property.createProperty("property1", "value1"),
+                        Property.createProperty("property2", "value2")),
+                    new SortedArrayStringMap())
+                .toMap())
         .containsExactly(
             "property1",
             "value1",
@@ -117,19 +126,19 @@ public final class OpenCensusTraceContextDataInjectorTest {
   @Test
   public void handleEmptyConfigurationProperties() {
     assertContainsOnlyDefaultTracingEntries(
-        new OpenCensusTraceContextDataInjector()
+        new OpenCensusTraceContextDataInjector(SpanSelection.ALL_SPANS)
             .injectContextData(Collections.<Property>emptyList(), new SortedArrayStringMap()));
   }
 
   @Test
   public void handleNullConfigurationProperties() {
     assertContainsOnlyDefaultTracingEntries(
-        new OpenCensusTraceContextDataInjector()
+        new OpenCensusTraceContextDataInjector(SpanSelection.ALL_SPANS)
             .injectContextData(null, new SortedArrayStringMap()));
   }
 
   private static void assertContainsOnlyDefaultTracingEntries(StringMap stringMap) {
-    assertThat(toMap(stringMap))
+    assertThat(stringMap.toMap())
         .containsExactly(
             "opencensusTraceId",
             "00000000000000000000000000000000",
@@ -139,15 +148,60 @@ public final class OpenCensusTraceContextDataInjectorTest {
             "false");
   }
 
-  private static Map<String, String> toMap(StringMap stringMap) {
-    final ImmutableMap.Builder<String, String> builder = ImmutableMap.builder();
-    stringMap.forEach(
-        new BiConsumer<String, Object>() {
-          @Override
-          public void accept(String key, Object value) {
-            builder.put(key, String.valueOf(value));
-          }
-        });
-    return builder.build();
+  @Test
+  public void rawContextDataWithTracingData() {
+    OpenCensusTraceContextDataInjector plugin =
+        new OpenCensusTraceContextDataInjector(SpanSelection.ALL_SPANS);
+    SpanContext spanContext =
+        SpanContext.create(
+            TraceId.fromLowerBase16("e17944156660f55b8cae5ce3f45d4a40"),
+            SpanId.fromLowerBase16("fc3d2ba0d283b66a"),
+            TraceOptions.builder().setIsSampled(true).build(),
+            EMPTY_TRACESTATE);
+    Scope scope = tracer.withSpan(new TestSpan(spanContext));
+    try {
+      String key = "myTestKey";
+      ThreadContext.put(key, "myTestValue");
+      try {
+        assertThat(plugin.rawContextData().toMap())
+            .containsExactly(
+                "myTestKey",
+                "myTestValue",
+                "opencensusTraceId",
+                "e17944156660f55b8cae5ce3f45d4a40",
+                "opencensusSpanId",
+                "fc3d2ba0d283b66a",
+                "opencensusTraceSampled",
+                "true");
+      } finally {
+        ThreadContext.remove(key);
+      }
+    } finally {
+      scope.close();
+    }
+  }
+
+  @Test
+  public void rawContextDataWithoutTracingData() {
+    OpenCensusTraceContextDataInjector plugin =
+        new OpenCensusTraceContextDataInjector(SpanSelection.NO_SPANS);
+    SpanContext spanContext =
+        SpanContext.create(
+            TraceId.fromLowerBase16("ea236000f6d387fe7c06c5a6d6458b53"),
+            SpanId.fromLowerBase16("f3b39dbbadb73074"),
+            TraceOptions.builder().setIsSampled(true).build(),
+            EMPTY_TRACESTATE);
+    Scope scope = tracer.withSpan(new TestSpan(spanContext));
+    try {
+      String key = "myTestKey";
+      ThreadContext.put(key, "myTestValue");
+      try {
+        assertThat(plugin.rawContextData().toMap()).containsExactly("myTestKey", "myTestValue");
+      } finally {
+        ThreadContext.remove(key);
+      }
+    } finally {
+      scope.close();
+    }
   }
 }
diff --git a/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/TestSpan.java b/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/TestSpan.java
new file mode 100644
index 00000000..7af46064
--- /dev/null
+++ b/contrib/log_correlation/log4j2/src/test/java/io/opencensus/contrib/logcorrelation/log4j2/TestSpan.java
@@ -0,0 +1,46 @@
+/*
+ * Copyright 2018, OpenCensus Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.opencensus.contrib.logcorrelation.log4j2;
+
+import io.opencensus.trace.Annotation;
+import io.opencensus.trace.AttributeValue;
+import io.opencensus.trace.EndSpanOptions;
+import io.opencensus.trace.Link;
+import io.opencensus.trace.Span;
+import io.opencensus.trace.SpanContext;
+import java.util.EnumSet;
+import java.util.Map;
+
+// Simple test Span that holds a SpanContext. The tests cannot use Span directly, since it is
+// abstract.
+final class TestSpan extends Span {
+  TestSpan(SpanContext context) {
+    super(context, EnumSet.of(Options.RECORD_EVENTS));
+  }
+
+  @Override
+  public void end(EndSpanOptions options) {}
+
+  @Override
+  public void addLink(Link link) {}
+
+  @Override
+  public void addAnnotation(Annotation annotation) {}
+
+  @Override
+  public void addAnnotation(String description, Map<String, AttributeValue> attributes) {}
+}