Adding unit test cases

Signed-off-by: Nishchay Malhotra <nishcha@amazon.com>

libs/telemetry/src/test/java/org/opensearch/telemetry/tracing/DefaultTracerTests.java

@@ -13,12 +13,17 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.node.Node;
 import org.opensearch.telemetry.tracing.attributes.Attributes;
+import org.opensearch.telemetry.tracing.attributes.SamplingAttributes;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.test.telemetry.tracing.MockSpan;
 import org.opensearch.test.telemetry.tracing.MockTracingTelemetry;
 import org.opensearch.threadpool.ThreadPool;
 
 import java.io.IOException;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 
@@ -424,4 +429,59 @@ private SpanCreationContext buildSpanCreationContext(String spanName, Attributes
         }
         return spanCreationContext;
     }
+
+    public void testCreateSpanWithInferredAttributes() {
+        TracingTelemetry tracingTelemetry = new MockTracingTelemetry();
+        ThreadContext threadContext = new ThreadContext(Settings.EMPTY);
+        DefaultTracer defaultTracer = new DefaultTracer(
+            tracingTelemetry,
+            new ThreadContextBasedTracerContextStorage(threadContext, tracingTelemetry)
+        );
+
+        SpanCreationContext spanCreationContext = buildSpanCreationContext(
+            "span_name",
+            Attributes.create().addAttribute(SamplingAttributes.SAMPLER.getValue(), SamplingAttributes.INFERRED_SAMPLER.getValue()),
+            null
+        );
+
+        Span span = defaultTracer.startSpan(spanCreationContext);
+
+        assertThat(defaultTracer.getCurrentSpan(), is(nullValue()));
+        assertEquals(SamplingAttributes.INFERRED_SAMPLER.getValue(), ((MockSpan) span).getAttribute(SamplingAttributes.SAMPLER.getValue()));
+        span.endSpan();
+    }
+
+    @SuppressWarnings("unchecked")
+    public void testCreateSpanWithInferredSampledParent() {
+        TracingTelemetry tracingTelemetry = new MockTracingTelemetry();
+        DefaultTracer defaultTracer = new DefaultTracer(
+            tracingTelemetry,
+            new ThreadContextBasedTracerContextStorage(new ThreadContext(Settings.EMPTY), tracingTelemetry)
+        );
+
+        SpanCreationContext spanCreationContext = buildSpanCreationContext("span_name", null, null);
+
+        Span span = defaultTracer.startSpan(
+            spanCreationContext,
+            (Map<String, Collection<String>>) new HashMap<String, Collection<String>>().put(
+                SamplingAttributes.SAMPLER.getValue(),
+                Collections.singleton(SamplingAttributes.INFERRED_SAMPLER.getValue())
+            )
+        );
+
+        try (final SpanScope scope = defaultTracer.withSpanInScope(span)) {
+            SpanContext parentSpan = defaultTracer.getCurrentSpan();
+            SpanCreationContext spanCreationContext1 = buildSpanCreationContext("span_name_1", Attributes.EMPTY, parentSpan.getSpan());
+            Span span1 = defaultTracer.startSpan(spanCreationContext1);
+            assertEquals("span_name_1", span1.getSpanName());
+            assertEquals(parentSpan.getSpan(), span1.getParentSpan());
+            assertEquals(
+                SamplingAttributes.INFERRED_SAMPLER.getValue(),
+                ((MockSpan) span1).getAttribute(SamplingAttributes.SAMPLER.getValue())
+            );
+            span1.endSpan();
+        } finally {
+            span.endSpan();
+        }
+    }
 }

plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelResourceProvider.java

@@ -12,7 +12,7 @@
 import org.opensearch.telemetry.TelemetrySettings;
 import org.opensearch.telemetry.metrics.exporter.OTelMetricsExporterFactory;
 import org.opensearch.telemetry.tracing.exporter.OTelSpanExporterFactory;
-import org.opensearch.telemetry.tracing.processor.OtelSpanProcessor;
+import org.opensearch.telemetry.tracing.processor.OTelSpanProcessor;
 import org.opensearch.telemetry.tracing.sampler.OTelSamplerFactory;
 import org.opensearch.telemetry.tracing.sampler.RequestSampler;
 
@@ -118,8 +118,8 @@ private static SdkTracerProvider createSdkTracerProvider(
             .build();
     }
 
-    private static OtelSpanProcessor spanProcessor(Settings settings, SpanExporter spanExporter) {
-        return new OtelSpanProcessor(batchSpanProcessor(settings, spanExporter));
+    private static OTelSpanProcessor spanProcessor(Settings settings, SpanExporter spanExporter) {
+        return new OTelSpanProcessor(batchSpanProcessor(settings, spanExporter));
     }
 
     private static BatchSpanProcessor batchSpanProcessor(Settings settings, SpanExporter spanExporter) {

plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/OtelSpanProcessor.java → plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessor.java

@@ -23,15 +23,15 @@
 /**
  * Implementation of the SpanProcessor and delegates to the configured processor.
  */
-public class OtelSpanProcessor implements SpanProcessor {
+public class OTelSpanProcessor implements SpanProcessor {
 
     private final SpanProcessor delegateProcessor;
 
     /**
      * *
      * @param delegateProcessor the span processor to which this processor will delegate
      */
-    public OtelSpanProcessor(SpanProcessor delegateProcessor) {
+    public OTelSpanProcessor(SpanProcessor delegateProcessor) {
         this.delegateProcessor = delegateProcessor;
     }
 

plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/OTelSpanTests.java

@@ -85,6 +85,30 @@ public void testGetSpanId() {
         assertEquals(SPAN_ID, oTelSpan.getSpanId());
     }
 
+    public void testGetSpanBoolean() {
+        Span mockSpan = getMockSpan();
+        OTelSpan oTelSpan = new OTelSpan("spanName", mockSpan, null);
+        assertNull(oTelSpan.getAttributeBoolean("key"));
+    }
+
+    public void testGetSpanString() {
+        Span mockSpan = getMockSpan();
+        OTelSpan oTelSpan = new OTelSpan("spanName", mockSpan, null);
+        assertNull(oTelSpan.getAttributeString("key"));
+    }
+
+    public void testGetSpanLong() {
+        Span mockSpan = getMockSpan();
+        OTelSpan oTelSpan = new OTelSpan("spanName", mockSpan, null);
+        assertNull(oTelSpan.getAttributeLong("key"));
+    }
+
+    public void testGetSpanDouble() {
+        Span mockSpan = getMockSpan();
+        OTelSpan oTelSpan = new OTelSpan("spanName", mockSpan, null);
+        assertNull(oTelSpan.getAttributeDouble("key"));
+    }
+
     private Span getMockSpan() {
         Span mockSpan = mock(Span.class);
         when(mockSpan.getSpanContext()).thenReturn(SpanContext.create(TRACE_ID, SPAN_ID, TraceFlags.getDefault(), TraceState.getDefault()));

plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/MockSpanProcessor.java

@@ -0,0 +1,47 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.telemetry.tracing.processor;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.common.CompletableResultCode;
+import io.opentelemetry.sdk.trace.ReadWriteSpan;
+import io.opentelemetry.sdk.trace.ReadableSpan;
+import io.opentelemetry.sdk.trace.SpanProcessor;
+
+public class MockSpanProcessor implements SpanProcessor {
+
+    @Override
+    public void onStart(Context parentContext, ReadWriteSpan span) {
+        {}
+    }
+
+    @Override
+    public boolean isStartRequired() {
+        return false;
+    }
+
+    @Override
+    public void onEnd(ReadableSpan span) {}
+
+    @Override
+    public boolean isEndRequired() {
+        return false;
+    }
+
+    @Override
+    public CompletableResultCode shutdown() {
+        return CompletableResultCode.ofSuccess();
+    }
+
+    @Override
+    public CompletableResultCode forceFlush() {
+        return CompletableResultCode.ofSuccess();
+    }
+
+}

plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessorTests.java

@@ -0,0 +1,84 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.telemetry.tracing.processor;
+
+import org.opensearch.telemetry.tracing.attributes.SamplingAttributes;
+import org.opensearch.test.OpenSearchTestCase;
+
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.TraceFlags;
+import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.sdk.trace.ReadableSpan;
+import io.opentelemetry.sdk.trace.SpanProcessor;
+import org.mockito.Mockito;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class OTelSpanProcessorTests extends OpenSearchTestCase {
+
+    private OTelSpanProcessor oTelSpanProcessor;
+    private MockSpanProcessor mockSpanProcessor;
+    private static final String TRACE_ID = "4aa59968f31dcbff7807741afa9d7d62";
+    private static final String SPAN_ID = "bea205cd25756b5e";
+
+    private ReadableSpan createEndedSpan() {
+        return Mockito.mock(ReadableSpan.class);
+    }
+
+    public void testOTelSpanProcessorDelegation() {
+        mockSpanProcessor = new MockSpanProcessor();
+        oTelSpanProcessor = new OTelSpanProcessor(mockSpanProcessor);
+        assertEquals(mockSpanProcessor.forceFlush(), oTelSpanProcessor.forceFlush());
+        assertEquals(mockSpanProcessor.shutdown(), oTelSpanProcessor.shutdown());
+        assertEquals(mockSpanProcessor.isEndRequired(), oTelSpanProcessor.isEndRequired());
+        assertEquals(mockSpanProcessor.isStartRequired(), oTelSpanProcessor.isStartRequired());
+    }
+
+    public void testOnEndFunction() {
+        SpanProcessor mockProcessor = mock(SpanProcessor.class);
+        oTelSpanProcessor = new OTelSpanProcessor(mockProcessor);
+        ReadableSpan readableSpan = this.createEndedSpan();
+        when(readableSpan.getSpanContext()).thenReturn(
+            SpanContext.create(TRACE_ID, SPAN_ID, TraceFlags.getDefault(), TraceState.getDefault())
+        );
+        oTelSpanProcessor.onEnd(readableSpan);
+        Mockito.verify(mockProcessor, Mockito.times(1)).onEnd(readableSpan);
+    }
+
+    public void testOnEndFunctionWithInferredAttribute() {
+        SpanProcessor mockProcessor = mock(SpanProcessor.class);
+        oTelSpanProcessor = new OTelSpanProcessor(mockProcessor);
+        ReadableSpan readableSpan = this.createEndedSpan();
+        when(readableSpan.getSpanContext()).thenReturn(
+            SpanContext.create(TRACE_ID, SPAN_ID, TraceFlags.getSampled(), TraceState.getDefault())
+        );
+        when(readableSpan.getAttribute(AttributeKey.stringKey(SamplingAttributes.SAMPLER.getValue()))).thenReturn(
+            SamplingAttributes.INFERRED_SAMPLER.getValue()
+        );
+        oTelSpanProcessor.onEnd(readableSpan);
+        Mockito.verify(mockProcessor, Mockito.times(0)).onEnd(readableSpan);
+    }
+
+    public void testOnEndFunctionWithInferredAttributeAndSampled() {
+        SpanProcessor mockProcessor = mock(SpanProcessor.class);
+        oTelSpanProcessor = new OTelSpanProcessor(mockProcessor);
+        ReadableSpan readableSpan = this.createEndedSpan();
+        when(readableSpan.getSpanContext()).thenReturn(
+            SpanContext.create(TRACE_ID, SPAN_ID, TraceFlags.getSampled(), TraceState.getDefault())
+        );
+        when(readableSpan.getAttribute(AttributeKey.stringKey(SamplingAttributes.SAMPLER.getValue()))).thenReturn(
+            SamplingAttributes.INFERRED_SAMPLER.getValue()
+        );
+        when(readableSpan.getAttribute(AttributeKey.booleanKey(SamplingAttributes.SAMPLED.getValue()))).thenReturn(true);
+        oTelSpanProcessor.onEnd(readableSpan);
+        Mockito.verify(mockProcessor, Mockito.times(1)).onEnd(readableSpan);
+    }
+}

plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/InferredSamplerTests.java

@@ -0,0 +1,109 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.telemetry.tracing.sampler;
+
+import org.opensearch.common.settings.ClusterSettings;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.telemetry.TelemetrySettings;
+import org.opensearch.telemetry.tracing.attributes.SamplingAttributes;
+import org.opensearch.test.OpenSearchTestCase;
+
+import java.util.Collections;
+import java.util.Set;
+
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.sdk.trace.samplers.Sampler;
+import io.opentelemetry.sdk.trace.samplers.SamplingResult;
+
+import static org.opensearch.telemetry.TelemetrySettings.TRACER_ENABLED_SETTING;
+import static org.opensearch.telemetry.TelemetrySettings.TRACER_INFERRED_SAMPLER_ALLOWLISTED;
+import static org.opensearch.telemetry.TelemetrySettings.TRACER_SAMPLER_PROBABILITY;
+import static org.opensearch.telemetry.tracing.AttributeNames.TRANSPORT_ACTION;
+import static org.mockito.Mockito.mock;
+
+public class InferredSamplerTests extends OpenSearchTestCase {
+
+    public void testGetSamplerWithSettingDisabled() {
+        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 inferredActionSampler = InferredActionSampler.create(telemetrySettings, Settings.EMPTY, null);
+
+        SamplingResult result = inferredActionSampler.shouldSample(
+            mock(Context.class),
+            "00000000000000000000000000000000",
+            "spanName",
+            SpanKind.INTERNAL,
+            Attributes.builder().put(TRANSPORT_ACTION, "dummy_action").build(),
+            Collections.emptyList()
+        );
+
+        assertEquals(SamplingResult.drop(), result);
+    }
+
+    public void testGetSamplerWithSettingEnabled() {
+        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);
+        telemetrySettings.setInferredSamplingAllowListed(true);
+
+        // InferredActionSampler
+        Sampler inferredActionSampler = InferredActionSampler.create(telemetrySettings, Settings.EMPTY, null);
+
+        SamplingResult result = inferredActionSampler.shouldSample(
+            mock(Context.class),
+            "00000000000000000000000000000000",
+            "spanName",
+            SpanKind.INTERNAL,
+            Attributes.builder().put(TRANSPORT_ACTION, "dummy_action").build(),
+            Collections.emptyList()
+        );
+
+        assertEquals(SamplingResult.recordAndSample().getDecision(), result.getDecision());
+    }
+
+    public void testGetSamplerWithAddedAttributes() {
+        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);
+        telemetrySettings.setInferredSamplingAllowListed(true);
+
+        // InferredActionSampler
+        Sampler inferredActionSampler = InferredActionSampler.create(telemetrySettings, Settings.EMPTY, null);
+
+        SamplingResult result = inferredActionSampler.shouldSample(
+            mock(Context.class),
+            "00000000000000000000000000000000",
+            "spanName",
+            SpanKind.INTERNAL,
+            Attributes.builder().put(TRANSPORT_ACTION, "dummy_action").build(),
+            Collections.emptyList()
+        );
+
+        assertTrue(result.getAttributes().asMap().containsKey(AttributeKey.stringKey(SamplingAttributes.SAMPLER.getValue())));
+        assertEquals(
+            result.getAttributes().get(AttributeKey.stringKey(SamplingAttributes.SAMPLER.getValue())),
+            SamplingAttributes.INFERRED_SAMPLER.getValue()
+        );
+    }
+}