Exporter/OcAgent: Add methods for getting and applying config. (#1488)

2 files changed, 61 insertions, 14 deletions

diff --git a/exporters/trace/ocagent/src/main/java/io/opencensus/exporter/trace/ocagent/TraceProtoUtils.java b/exporters/trace/ocagent/src/main/java/io/opencensus/exporter/trace/ocagent/TraceProtoUtils.java
index df11ef58..ec778ba6 100644
--- a/exporters/trace/ocagent/src/main/java/io/opencensus/exporter/trace/ocagent/TraceProtoUtils.java
+++ b/exporters/trace/ocagent/src/main/java/io/opencensus/exporter/trace/ocagent/TraceProtoUtils.java
@@ -23,6 +23,7 @@ import com.google.protobuf.UInt32Value;
 import io.opencensus.common.Function;
 import io.opencensus.common.Functions;
 import io.opencensus.common.Timestamp;
+import io.opencensus.proto.agent.trace.v1.UpdatedLibraryConfig;
 import io.opencensus.proto.trace.v1.AttributeValue;
 import io.opencensus.proto.trace.v1.ConstantSampler;
 import io.opencensus.proto.trace.v1.ProbabilitySampler;
@@ -58,12 +59,8 @@ import java.util.Map;
 import org.checkerframework.checker.nullness.qual.Nullable;
 */
 
-/**
- * Utilities for converting the Tracing data models in OpenCensus Java to/from OpenCensus Proto.
- *
- * @since 0.17
- */
-public final class TraceProtoUtils {
+/** Utilities for converting the Tracing data models in OpenCensus Java to/from OpenCensus Proto. */
+final class TraceProtoUtils {
 
   // Constant functions for AttributeValue.
   private static final Function<String, /*@Nullable*/ AttributeValue> stringAttributeValueFunction =
@@ -106,10 +103,8 @@ public final class TraceProtoUtils {
    *
    * @param spanData the {@code SpanData}.
    * @return proto representation of {@code Span}.
-   * @since 0.17
    */
-  @SuppressWarnings("DefaultCharset")
-  public static Span toSpanProto(SpanData spanData) {
+  static Span toSpanProto(SpanData spanData) {
     SpanContext spanContext = spanData.getContext();
     TraceId traceId = spanContext.getTraceId();
     SpanId spanId = spanContext.getSpanId();
@@ -298,7 +293,6 @@ public final class TraceProtoUtils {
     }
   }
 
-  @SuppressWarnings("DefaultCharset")
   private static Link toLinkProto(io.opencensus.trace.Link link) {
     return Link.newBuilder()
         .setTraceId(toByteString(link.getTraceId().getBytes()))
@@ -322,9 +316,8 @@ public final class TraceProtoUtils {
    *
    * @param traceParams the {@code TraceParams}.
    * @return {@code TraceConfig}.
-   * @since 0.17
    */
-  public static TraceConfig toTraceConfigProto(TraceParams traceParams) {
+  static TraceConfig toTraceConfigProto(TraceParams traceParams) {
     TraceConfig.Builder traceConfigProtoBuilder = TraceConfig.newBuilder();
     Sampler librarySampler = traceParams.getSampler();
 
@@ -360,7 +353,7 @@ public final class TraceProtoUtils {
    * @return updated {@code TraceParams}.
    * @since 0.17
    */
-  public static TraceParams fromTraceConfigProto(
+  static TraceParams fromTraceConfigProto(
       TraceConfig traceConfigProto, TraceParams currentTraceParams) {
     TraceParams.Builder builder = currentTraceParams.toBuilder();
     if (traceConfigProto.hasConstantSampler()) {
@@ -378,5 +371,20 @@ public final class TraceProtoUtils {
     return builder.build();
   }
 
+  // Creates a TraceConfig proto message with current TraceParams.
+  static TraceConfig getCurrentTraceConfig(io.opencensus.trace.config.TraceConfig traceConfig) {
+    TraceParams traceParams = traceConfig.getActiveTraceParams();
+    return toTraceConfigProto(traceParams);
+  }
+
+  // Creates an updated TraceParams with the given UpdatedLibraryConfig message and current
+  // TraceParams, then applies the updated TraceParams.
+  static TraceParams getUpdatedTraceParams(
+      UpdatedLibraryConfig config, io.opencensus.trace.config.TraceConfig traceConfig) {
+    TraceParams currentParams = traceConfig.getActiveTraceParams();
+    TraceConfig traceConfigProto = config.getConfig();
+    return fromTraceConfigProto(traceConfigProto, currentParams);
+  }
+
   private TraceProtoUtils() {}
 }
diff --git a/exporters/trace/ocagent/src/test/java/io/opencensus/exporter/trace/ocagent/TraceProtoUtilsTest.java b/exporters/trace/ocagent/src/test/java/io/opencensus/exporter/trace/ocagent/TraceProtoUtilsTest.java
index 4271fc79..74c7c29e 100644
--- a/exporters/trace/ocagent/src/test/java/io/opencensus/exporter/trace/ocagent/TraceProtoUtilsTest.java
+++ b/exporters/trace/ocagent/src/test/java/io/opencensus/exporter/trace/ocagent/TraceProtoUtilsTest.java
@@ -25,6 +25,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.protobuf.BoolValue;
 import com.google.protobuf.UInt32Value;
 import io.opencensus.common.Timestamp;
+import io.opencensus.proto.agent.trace.v1.UpdatedLibraryConfig;
 import io.opencensus.proto.trace.v1.AttributeValue;
 import io.opencensus.proto.trace.v1.ConstantSampler;
 import io.opencensus.proto.trace.v1.ProbabilitySampler;
@@ -49,14 +50,20 @@ import io.opencensus.trace.export.SpanData.TimedEvent;
 import io.opencensus.trace.export.SpanData.TimedEvents;
 import io.opencensus.trace.samplers.Samplers;
 import java.util.List;
+import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
 
 /** Tests for {@link TraceProtoUtils}. */
 @RunWith(JUnit4.class)
 public class TraceProtoUtilsTest {
 
+  @Mock private io.opencensus.trace.config.TraceConfig mockTraceConfig;
+
   private static final TraceParams DEFAULT_PARAMS = TraceParams.DEFAULT;
 
   private static final Timestamp startTimestamp = Timestamp.create(123, 456);
@@ -126,7 +133,15 @@ public class TraceProtoUtilsTest {
       TimedEvents.create(networkEventsList, DROPPED_NETWORKEVENTS_COUNT);
   private static final SpanData.Links links = SpanData.Links.create(linksList, DROPPED_LINKS_COUNT);
 
-  @SuppressWarnings("DefaultCharset")
+  @Before
+  public void setUp() {
+    MockitoAnnotations.initMocks(this);
+    Mockito.when(mockTraceConfig.getActiveTraceParams()).thenReturn(DEFAULT_PARAMS);
+    Mockito.doNothing()
+        .when(mockTraceConfig)
+        .updateActiveTraceParams(Mockito.any(TraceParams.class));
+  }
+
   @Test
   public void toSpanProto() {
     SpanData spanData =
@@ -312,6 +327,30 @@ public class TraceProtoUtilsTest {
         .isEqualTo(Samplers.probabilitySampler(0.01));
   }
 
+  @Test
+  public void getCurrentTraceConfig() {
+    TraceConfig configProto = TraceProtoUtils.toTraceConfigProto(DEFAULT_PARAMS);
+    assertThat(TraceProtoUtils.getCurrentTraceConfig(mockTraceConfig)).isEqualTo(configProto);
+    Mockito.verify(mockTraceConfig, Mockito.times(1)).getActiveTraceParams();
+  }
+
+  @Test
+  public void applyUpdatedConfig() {
+    TraceConfig configProto =
+        TraceConfig.newBuilder()
+            .setProbabilitySampler(
+                ProbabilitySampler.newBuilder().setSamplingProbability(0.01).build())
+            .build();
+    UpdatedLibraryConfig updatedLibraryConfig =
+        UpdatedLibraryConfig.newBuilder().setConfig(configProto).build();
+    TraceParams traceParams =
+        TraceProtoUtils.getUpdatedTraceParams(updatedLibraryConfig, mockTraceConfig);
+    TraceParams expectedParams =
+        DEFAULT_PARAMS.toBuilder().setSampler(Samplers.probabilitySampler(0.01)).build();
+    Mockito.verify(mockTraceConfig, Mockito.times(1)).getActiveTraceParams();
+    assertThat(traceParams).isEqualTo(expectedParams);
+  }
+
   private static TraceParams getTraceParams(Sampler sampler) {
     return DEFAULT_PARAMS.toBuilder().setSampler(sampler).build();
   }