Add logging handler test for traceEnhancer

Author: cindy-peng

google-cloud-logging/src/main/java/com/google/cloud/logging/TraceLoggingEnhancer.java

@@ -25,6 +25,7 @@ public TraceLoggingEnhancer(String prefix) {}
 
   private static final ThreadLocal<String> traceId = new ThreadLocal<>();
   private static final ThreadLocal<String> spanId = new ThreadLocal<>();
+  private static final ThreadLocal<Boolean> traceSampled = new ThreadLocal<Boolean>();
 
   /**
    * Set the Trace ID associated with any logging done by the current thread.
@@ -52,6 +53,19 @@ public static void setCurrentSpanId(String id) {
     }
   }
 
+  /**
+   * Set the trace sampled flag associated with any logging done by the current thread.
+   *
+   * @param isTraceSampled The traceSampled flag
+   */
+  public static void setCurrentTraceSampled(Boolean isTraceSampled) {
+    if (isTraceSampled == null) {
+      traceSampled.remove();
+    } else {
+      traceSampled.set(isTraceSampled);
+    }
+  }
+
   /**
    * Get the Trace ID associated with any logging done by the current thread.
    *
@@ -70,6 +84,15 @@ public static String getCurrentSpanId() {
     return spanId.get();
   }
 
+  /**
+   * Get the trace sampled flag associated with any logging done by the current thread.
+   *
+   * @return traceSampled The traceSampled flag
+   */
+  public static Boolean getCurrentTraceSampled() {
+    return traceSampled.get();
+  }
+
   @Override
   public void enhanceLogEntry(LogEntry.Builder builder) {
     String traceId = getCurrentTraceId();
@@ -80,5 +103,9 @@ public void enhanceLogEntry(LogEntry.Builder builder) {
     if (spanId != null) {
       builder.setSpanId(spanId);
     }
+    Boolean isTraceSampled = getCurrentTraceSampled();
+    if (isTraceSampled != null) {
+      builder.setTraceSampled(isTraceSampled);
+    }
   }
 }

google-cloud-logging/src/test/java/com/google/cloud/logging/LoggingHandlerTest.java

@@ -165,6 +165,8 @@ public class LoggingHandlerTest {
           .addLabel("levelName", "FINEST")
           .addLabel("levelValue", String.valueOf(Level.FINEST.intValue()))
           .setTrace("projects/projectId/traces/traceId")
+          .setSpanId("test_span_id")
+          .setTraceSampled(true)
           .setTimestamp(123456789L)
           .build();
   private static final LogEntry DIAGNOSTIC_ENTRY =
@@ -454,6 +456,8 @@ public void testTraceEnhancedLogEntry() {
     replay(options, logging);
     LoggingEnhancer enhancer = new TraceLoggingEnhancer();
     TraceLoggingEnhancer.setCurrentTraceId("projects/projectId/traces/traceId");
+    TraceLoggingEnhancer.setCurrentSpanId("test_span_id");
+    TraceLoggingEnhancer.setCurrentTraceSampled(true);
     Handler handler =
         new LoggingHandler(LOG_NAME, options, DEFAULT_RESOURCE, ImmutableList.of(enhancer));
     handler.setLevel(Level.ALL);