From 5762b130f77cc1fa55cba354ec0c1b267450afcc Mon Sep 17 00:00:00 2001
From: jkwatson <jkwatson@gmail.com>
Date: Wed, 6 Jan 2021 14:20:16 -0800
Subject: [PATCH] Add a test for the RECORD_ONLY span And remove the old TODOs
 from the BSP test.

---
 .../trace/export/BatchSpanProcessorTest.java  | 82 ++++++++++---------
 1 file changed, 42 insertions(+), 40 deletions(-)

diff --git a/sdk/trace/src/test/java/io/opentelemetry/sdk/trace/export/BatchSpanProcessorTest.java b/sdk/trace/src/test/java/io/opentelemetry/sdk/trace/export/BatchSpanProcessorTest.java
index d1de790b033..1492e72a43c 100644
--- a/sdk/trace/src/test/java/io/opentelemetry/sdk/trace/export/BatchSpanProcessorTest.java
+++ b/sdk/trace/src/test/java/io/opentelemetry/sdk/trace/export/BatchSpanProcessorTest.java
@@ -11,14 +11,17 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Context;
 import io.opentelemetry.sdk.common.CompletableResultCode;
 import io.opentelemetry.sdk.trace.ReadableSpan;
 import io.opentelemetry.sdk.trace.SdkTracerProvider;
 import io.opentelemetry.sdk.trace.config.TraceConfig;
 import io.opentelemetry.sdk.trace.data.SpanData;
 import io.opentelemetry.sdk.trace.samplers.Sampler;
+import io.opentelemetry.sdk.trace.samplers.SamplingResult;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -50,16 +53,6 @@ void cleanup() {
     }
   }
 
-  // TODO(bdrutu): Fix this when Sampler return RECORD_ONLY option.
-  /*
-  private ReadableSpan createNotSampledRecordingEventsEndedSpan(String spanName) {
-    io.opentelemetry.trace.Span span =
-        tracer.spanBuilder(spanName).setSampler(Samplers.neverSample()).startSpanWithSampler();
-    span.end();
-    return (ReadableSpan) span;
-  }
-  */
-
   private ReadableSpan createEndedSpan(String spanName) {
     Tracer tracer = sdkTracerProvider.get(getClass().getName());
     Span span = tracer.spanBuilder(spanName).startSpan();
@@ -378,7 +371,7 @@ void exportNotSampledSpans() {
     sdkTracerProvider.get("test").spanBuilder(SPAN_NAME_1).startSpan().end();
     sdkTracerProvider.get("test").spanBuilder(SPAN_NAME_2).startSpan().end();
     traceConfig.set(traceConfig.get().toBuilder().setSampler(Sampler.alwaysOn()).build());
-    ReadableSpan span2 = createEndedSpan(SPAN_NAME_2);
+    ReadableSpan span = createEndedSpan(SPAN_NAME_2);
     // Spans are recorded and exported in the same order as they are ended, we test that a non
     // sampled span is not exported by creating and ending a sampled span after a non sampled span
     // and checking that the first exported span is the sampled span (the non sampled did not get
@@ -386,40 +379,49 @@ void exportNotSampledSpans() {
     List<SpanData> exported = waitingSpanExporter.waitForExport();
     // Need to check this because otherwise the variable span1 is unused, other option is to not
     // have a span1 variable.
-    assertThat(exported).containsExactly(span2.toSpanData());
+    assertThat(exported).containsExactly(span.toSpanData());
   }
 
   @Test
-  void exportNotSampledSpans_recordingEvents() {
-    // TODO(bdrutu): Fix this when Sampler return RECORD_ONLY option.
-    /*
-    sdkTracerProvider.addSpanProcessor(
-        BatchSpanProcessor.builder(waitingSpanExporter)
-            .setScheduleDelayMillis(MAX_SCHEDULE_DELAY_MILLIS)
-            .reportOnlySampled(false)
-            .build());
+  void exportNotSampledSpans_recordOnly() {
+    WaitingSpanExporter waitingSpanExporter =
+        new WaitingSpanExporter(1, CompletableResultCode.ofSuccess());
+    AtomicReference<TraceConfig> traceConfig =
+        new AtomicReference<>(
+            TraceConfig.getDefault().toBuilder()
+                .setSampler(
+                    new Sampler() {
+                      @Override
+                      public SamplingResult shouldSample(
+                          Context parentContext,
+                          String traceId,
+                          String name,
+                          Span.Kind spanKind,
+                          Attributes attributes,
+                          List<SpanData.Link> parentLinks) {
+                        return SamplingResult.create(SamplingResult.Decision.RECORD_ONLY);
+                      }
 
-    ReadableSpan span = createNotSampledRecordingEventsEndedSpan(SPAN_NAME_1);
-    List<SpanData> exported = waitingSpanExporter.waitForExport(1);
-    assertThat(exported).containsExactly(span.toSpanData());
-    */
-  }
+                      @Override
+                      public String getDescription() {
+                        return null;
+                      }
+                    })
+                .build());
+    sdkTracerProvider =
+        SdkTracerProvider.builder()
+            .addSpanProcessor(
+                BatchSpanProcessor.builder(waitingSpanExporter)
+                    .setScheduleDelayMillis(MAX_SCHEDULE_DELAY_MILLIS)
+                    .setExportOnlySampled(true)
+                    .build())
+            .setTraceConfig(traceConfig::get)
+            .build();
 
-  @Test
-  void exportNotSampledSpans_reportOnlySampled() {
-    // TODO(bdrutu): Fix this when Sampler return RECORD_ONLY option.
-    /*
-    sdkTracerProvider.addSpanProcessor(
-        BatchSpanProcessor.builder(waitingSpanExporter)
-            .reportOnlySampled(true)
-            .setScheduleDelayMillis(MAX_SCHEDULE_DELAY_MILLIS)
-            .build());
-
-    createNotSampledRecordingEventsEndedSpan(SPAN_NAME_1);
-    ReadableSpan sampledSpan = createSampledEndedSpan(SPAN_NAME_2);
-    List<SpanData> exported = waitingSpanExporter.waitForExport(1);
-    assertThat(exported).containsExactly(sampledSpan.toSpanData());
-    */
+    createEndedSpan(SPAN_NAME_2);
+
+    List<SpanData> exported = waitingSpanExporter.waitForExport();
+    assertThat(exported).isEmpty();
   }
 
   @Test