Addimh more unit test cases

Signed-off-by: Nishchay Malhotra <nishcha@amazon.com>
opensearch-project · Apr 1, 2024 · b944c0b · b944c0b
1 parent 8bce540
commit b944c0b
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 1 deletion.
diff --git a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelSpan.java b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelSpan.java
@@ -10,6 +10,8 @@
 
 import org.opensearch.telemetry.tracing.attributes.SamplingAttributes;
 
+import java.util.Optional;
+
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.StatusCode;
@@ -47,7 +49,12 @@ public void endSpan() {
      * This Framework will be used to evaluate a span if that is an outlier or not.
      */
     private boolean isSpanOutlier() {
-        return false;
+        Optional<Boolean> isSpanSampled = Optional.ofNullable(getAttributeBoolean(SamplingAttributes.SAMPLED.getValue()));
+        Optional<String> isSpanInferredSampled = Optional.ofNullable(getAttributeString(SamplingAttributes.SAMPLER.getValue()));
+
+        return isSpanSampled.isPresent()
+            && isSpanInferredSampled.isPresent()
+            && isSpanInferredSampled.get().equals(SamplingAttributes.INFERRED_SAMPLER.getValue());
     }
 
     private void markParentForSampling() {

diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/OTelSpanTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/OTelSpanTests.java
@@ -8,12 +8,15 @@
 
 package org.opensearch.telemetry.tracing;
 
+import org.opensearch.telemetry.tracing.attributes.SamplingAttributes;
 import org.opensearch.test.OpenSearchTestCase;
 
+import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.TraceFlags;
 import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.sdk.trace.ReadWriteSpan;
 
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -109,6 +112,33 @@ public void testGetSpanDouble() {
         assertNull(oTelSpan.getAttributeDouble("key"));
     }
 
+    public void testSpanOutlier() {
+        ReadWriteSpan mockReadWriteSpan = mock(ReadWriteSpan.class);
+        Span mockSpan = getMockSpan();
+        OTelSpan mockParent = new OTelSpan("parentSpan", mockSpan, null);
+        OTelSpan oTelSpan = new OTelSpan("spanName", mockReadWriteSpan, mockParent);
+        when(mockReadWriteSpan.getAttribute(AttributeKey.booleanKey(SamplingAttributes.SAMPLED.getValue()))).thenReturn(true);
+        when(mockReadWriteSpan.getAttribute(AttributeKey.stringKey(SamplingAttributes.SAMPLER.getValue()))).thenReturn(
+            SamplingAttributes.INFERRED_SAMPLER.getValue()
+        );
+        oTelSpan.endSpan();
+        verify(mockSpan).setAttribute(SamplingAttributes.SAMPLED.getValue(), true);
+    }
+
+    public void testSpanAttributes() {
+        ReadWriteSpan mockReadWriteSpan = mock(ReadWriteSpan.class);
+        OTelSpan oTelSpan = new OTelSpan("spanName", mockReadWriteSpan, null);
+        oTelSpan.addAttribute("key", 0.0);
+        when(mockReadWriteSpan.getAttribute(AttributeKey.doubleKey("key"))).thenReturn(0.0);
+        verify(mockReadWriteSpan).setAttribute("key", 0.0);
+        assertEquals(0.0, (Object) oTelSpan.getAttributeDouble("key"));
+
+        oTelSpan.addAttribute("key1", 0L);
+        when(mockReadWriteSpan.getAttribute(AttributeKey.longKey("key1"))).thenReturn(0L);
+        verify(mockReadWriteSpan).setAttribute("key1", 0L);
+        assertEquals(0L, (Object) oTelSpan.getAttributeLong("key1"));
+    }
+
     private Span getMockSpan() {
         Span mockSpan = mock(Span.class);
         when(mockSpan.getSpanContext()).thenReturn(SpanContext.create(TRACE_ID, SPAN_ID, TraceFlags.getDefault(), TraceState.getDefault()));

diff --git a/...otel/src/test/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessorTests.java b/...otel/src/test/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessorTests.java
@@ -15,6 +15,8 @@
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.TraceFlags;
 import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.trace.ReadWriteSpan;
 import io.opentelemetry.sdk.trace.ReadableSpan;
 import io.opentelemetry.sdk.trace.SpanProcessor;
 import org.mockito.Mockito;
@@ -53,6 +55,15 @@ public void testOnEndFunction() {
         Mockito.verify(mockProcessor, Mockito.times(1)).onEnd(readableSpan);
     }
 
+    public void testOnStartFunction() {
+        SpanProcessor mockProcessor = mock(SpanProcessor.class);
+        Context spanContext = mock(Context.class);
+        oTelSpanProcessor = new OTelSpanProcessor(mockProcessor);
+        ReadWriteSpan readWriteSpan = mock(ReadWriteSpan.class);
+        oTelSpanProcessor.onStart(spanContext, readWriteSpan);
+        Mockito.verify(mockProcessor, Mockito.times(1)).onStart(spanContext, readWriteSpan);
+    }
+
     public void testOnEndFunctionWithInferredAttribute() {
         SpanProcessor mockProcessor = mock(SpanProcessor.class);
         oTelSpanProcessor = new OTelSpanProcessor(mockProcessor);

diff --git a/...try-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/InferredSamplerTests.java b/...try-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/InferredSamplerTests.java
@@ -105,5 +105,36 @@ public void testGetSamplerWithAddedAttributes() {
             result.getAttributes().get(AttributeKey.stringKey(SamplingAttributes.SAMPLER.getValue())),
             SamplingAttributes.INFERRED_SAMPLER.getValue()
         );
+        assertEquals("Inferred Action Sampler", inferredActionSampler.getDescription());
+    }
+
+    public void testFallBackSampler() {
+        ClusterSettings clusterSettings = new ClusterSettings(
+            Settings.EMPTY,
+            Set.of(TRACER_SAMPLER_PROBABILITY, TRACER_ENABLED_SETTING, TRACER_INFERRED_SAMPLER_ALLOWLISTED)
+        );
+
+        TelemetrySettings telemetrySettings = new TelemetrySettings(Settings.EMPTY, clusterSettings);
+
+        // InferredActionSampler
+        Sampler probabilisticTransportActionSampler = ProbabilisticTransportActionSampler.create(telemetrySettings, Settings.EMPTY, null);
+        Sampler inferredActionSampler = InferredActionSampler.create(
+            telemetrySettings,
+            Settings.EMPTY,
+            probabilisticTransportActionSampler
+        );
+
+        SamplingResult result = inferredActionSampler.shouldSample(
+            mock(Context.class),
+            "00000000000000000000000000000000",
+            "spanName",
+            SpanKind.INTERNAL,
+            Attributes.builder().put(TRANSPORT_ACTION, "dummy_action").build(),
+            Collections.emptyList()
+        );
+
+        // ProbabilisticTransportActionSampler
+        assertEquals(SamplingResult.recordAndSample(), result);
+        assertEquals(0.001, ((ProbabilisticTransportActionSampler) probabilisticTransportActionSampler).getSamplingRatio(), 0.000d);
     }
 }