From 8bce54060321fc31204a1a251e13298143694fe5 Mon Sep 17 00:00:00 2001 From: Nishchay Malhotra Date: Mon, 1 Apr 2024 16:58:33 +0530 Subject: [PATCH] Adding unit test cases Signed-off-by: Nishchay Malhotra --- .../telemetry/tracing/DefaultTracerTests.java | 60 ++++++++++ .../tracing/OTelResourceProvider.java | 6 +- ...nProcessor.java => OTelSpanProcessor.java} | 4 +- .../telemetry/tracing/OTelSpanTests.java | 24 ++++ .../tracing/processor/MockSpanProcessor.java | 47 ++++++++ .../processor/OTelSpanProcessorTests.java | 84 ++++++++++++++ .../tracing/sampler/InferredSamplerTests.java | 109 ++++++++++++++++++ .../OTelSamplingResultTests.java | 86 ++++++++++++++ .../telemetry/TelemetrySettingsTests.java | 30 +++++ .../telemetry/tracing/SpanBuilderTests.java | 14 +++ 10 files changed, 459 insertions(+), 5 deletions(-) rename plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/{OtelSpanProcessor.java => OTelSpanProcessor.java} (96%) create mode 100644 plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/MockSpanProcessor.java create mode 100644 plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessorTests.java create mode 100644 plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/InferredSamplerTests.java create mode 100644 plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/samplingResult/OTelSamplingResultTests.java diff --git a/libs/telemetry/src/test/java/org/opensearch/telemetry/tracing/DefaultTracerTests.java b/libs/telemetry/src/test/java/org/opensearch/telemetry/tracing/DefaultTracerTests.java index 2182b3ea28ac8..f088cd6054e76 100644 --- a/libs/telemetry/src/test/java/org/opensearch/telemetry/tracing/DefaultTracerTests.java +++ b/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>) new HashMap>().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(); + } + } } diff --git a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelResourceProvider.java b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelResourceProvider.java index d461a24081184..51150662ff9c6 100644 --- a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelResourceProvider.java +++ b/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) { diff --git a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/OtelSpanProcessor.java b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessor.java similarity index 96% rename from plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/OtelSpanProcessor.java rename to plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessor.java index d16f41e7ddd20..42d302458bf4b 100644 --- a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/OtelSpanProcessor.java +++ b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessor.java @@ -23,7 +23,7 @@ /** * Implementation of the SpanProcessor and delegates to the configured processor. */ -public class OtelSpanProcessor implements SpanProcessor { +public class OTelSpanProcessor implements SpanProcessor { private final SpanProcessor delegateProcessor; @@ -31,7 +31,7 @@ public class OtelSpanProcessor implements SpanProcessor { * * * @param delegateProcessor the span processor to which this processor will delegate */ - public OtelSpanProcessor(SpanProcessor delegateProcessor) { + public OTelSpanProcessor(SpanProcessor delegateProcessor) { this.delegateProcessor = delegateProcessor; } 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 index fc92ab36908e1..227c53f3ea4fb 100644 --- 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 @@ -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())); diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/MockSpanProcessor.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/MockSpanProcessor.java new file mode 100644 index 0000000000000..fc165cc9b5d01 --- /dev/null +++ b/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(); + } + +} diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessorTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/processor/OTelSpanProcessorTests.java new file mode 100644 index 0000000000000..09f1cf14fc0e2 --- /dev/null +++ b/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); + } +} diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/InferredSamplerTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/sampler/InferredSamplerTests.java new file mode 100644 index 0000000000000..6b3ca318f5572 --- /dev/null +++ b/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() + ); + } +} diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/samplingResult/OTelSamplingResultTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/samplingResult/OTelSamplingResultTests.java new file mode 100644 index 0000000000000..b85357532cc55 --- /dev/null +++ b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/samplingResult/OTelSamplingResultTests.java @@ -0,0 +1,86 @@ +/* + * 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.samplingResult; + +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.telemetry.tracing.sampler.InferredActionSampler; +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 OTelSamplingResultTests extends OpenSearchTestCase { + + public void testSamplingResult() { + 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.getClass(), OTelSamplingResult.class); + } + + public void testSamplingDecisionAndAttributes() { + 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() + ); + + OTelSamplingResult oTelSamplingResult = new OTelSamplingResult(result.getDecision(), result.getAttributes()); + assertEquals(result.getDecision(), oTelSamplingResult.getDecision()); + assertEquals(result.getAttributes(), oTelSamplingResult.getAttributes()); + } +} diff --git a/server/src/test/java/org/opensearch/telemetry/TelemetrySettingsTests.java b/server/src/test/java/org/opensearch/telemetry/TelemetrySettingsTests.java index 98af6ec3ea6d9..badc8d92345ac 100644 --- a/server/src/test/java/org/opensearch/telemetry/TelemetrySettingsTests.java +++ b/server/src/test/java/org/opensearch/telemetry/TelemetrySettingsTests.java @@ -71,4 +71,34 @@ public void testGetSamplingProbability() { assertEquals(0.02, telemetrySettings.getSamplingProbability(), 0.00d); } + public void testGetInferredSetting() { + 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); + + assertFalse(telemetrySettings.getInferredSamplingAllowListed()); + + clusterSettings.applySettings(Settings.builder().put("telemetry.inferred.sampler.allowlisted", "true").build()); + + // Validate inferred allowlist setting + assertTrue(telemetrySettings.getInferredSamplingAllowListed()); + } + + public void testSetInferredSetting() { + 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); + + assertFalse(telemetrySettings.getInferredSamplingAllowListed()); + + telemetrySettings.setInferredSamplingAllowListed(true); + + // Validate inferred allowlist setting + assertTrue(telemetrySettings.getInferredSamplingAllowListed()); + } + } diff --git a/server/src/test/java/org/opensearch/telemetry/tracing/SpanBuilderTests.java b/server/src/test/java/org/opensearch/telemetry/tracing/SpanBuilderTests.java index 4b763e4bd4454..96ab2c9399095 100644 --- a/server/src/test/java/org/opensearch/telemetry/tracing/SpanBuilderTests.java +++ b/server/src/test/java/org/opensearch/telemetry/tracing/SpanBuilderTests.java @@ -115,6 +115,20 @@ public void testParentSpan() { assertEquals(parentSpanContext, context.getParent()); } + public void testNoopSpanAttributes() { + String spanName = "test-name"; + SpanContext parentSpanContext = new SpanContext(NoopSpan.INSTANCE); + SpanCreationContext context = SpanBuilder.from(spanName, parentSpanContext); + Attributes attributes = context.getAttributes(); + assertNull(attributes); + assertEquals(spanName, context.getSpanName()); + assertEquals(parentSpanContext, context.getParent()); + assertEquals("", parentSpanContext.getSpan().getAttributeString("mockKey")); + assertEquals(false, parentSpanContext.getSpan().getAttributeBoolean("mockKey")); + assertEquals(0L, (Object) parentSpanContext.getSpan().getAttributeLong("mockKey")); + assertEquals(0.0, (Object) parentSpanContext.getSpan().getAttributeDouble("mockKey")); + } + private static Transport.Connection createTransportConnection() { return new Transport.Connection() { @Override