Skip to content

Commit

Permalink
add unit tests for def _custom_import_sampler()
Browse files Browse the repository at this point in the history
  • Loading branch information
jj22ee committed Feb 16, 2024
1 parent 6d1aa78 commit ba62648
Show file tree
Hide file tree
Showing 2 changed files with 103 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -106,11 +106,6 @@ def _initialize_components(auto_instrumentation_version):
sampler_name = _get_sampler()
sampler = _custom_import_sampler(sampler_name, resource)

if sampler is not None:
_logger.debug("Using sampler: %s", sampler.get_description())
else:
_logger.debug("Sampler is None")

_init_tracing(
exporters=trace_exporters,
id_generator=id_generator,
Expand Down Expand Up @@ -152,6 +147,8 @@ def _custom_import_sampler(sampler_name, resource) -> Sampler:
# TODO: Remove `sampler_name is None` condition when xray sampler is configured here:
# https://github.com/aws/amazon-cloudwatch-agent-operator/blob/main/pkg/instrumentation/defaultinstrumentation.go#L90
if sampler_name is None or sampler_name == "xray":
# Example env var value
# OTEL_TRACES_SAMPLER_ARG=endpoint=http://localhost:2000,polling_interval=360
sampler_argument_env = os.getenv(OTEL_TRACES_SAMPLER_ARG, None)
endpoint = None
polling_interval = None
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,10 @@
import time
from unittest import TestCase

from amazon.opentelemetry.distro.aws_opentelemetry_configurator import AwsOpenTelemetryConfigurator
from amazon.opentelemetry.distro.aws_opentelemetry_configurator import (
AwsOpenTelemetryConfigurator,
_custom_import_sampler,
)
from amazon.opentelemetry.distro.aws_opentelemetry_distro import AwsOpenTelemetryDistro
from opentelemetry.environment_variables import OTEL_LOGS_EXPORTER, OTEL_METRICS_EXPORTER, OTEL_TRACES_EXPORTER
from opentelemetry.sdk.environment_variables import OTEL_TRACES_SAMPLER, OTEL_TRACES_SAMPLER_ARG
Expand Down Expand Up @@ -52,3 +55,100 @@ def test_trace_id_ratio_sampler(self):
span.end()
# Configured for 1%, confirm there are at most 5% to account for randomness and reduce test flakiness.
self.assertGreater(0.05, num_sampled / num_spans)

# Test method for import xray
# Cannot test this logic via `aws_otel_configurator.configure()` because that will
# attempt to setup tracer provider again, which may cause issue
def test_import_xray_sampler_without_environment_arguments(self):
os.environ.pop(OTEL_TRACES_SAMPLER_ARG, None)

# May log http request error as xray sampler will attempt to fetch rules
xray_sampler = _custom_import_sampler(None, resource=None)
xray_sampler._rules_timer.cancel()
xray_sampler._targets_timer.cancel()
xray_client = xray_sampler._AwsXRayRemoteSampler__xray_client
self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 300)
self.assertEqual(
xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint, "http://127.0.0.1:2000/GetSamplingRules"
)

def test_import_xray_sampler_with_valid_environment_arguments(self):
os.environ.pop(OTEL_TRACES_SAMPLER_ARG, None)
os.environ.setdefault(OTEL_TRACES_SAMPLER_ARG, "endpoint=http://localhost:2000,polling_interval=600")

# May log http request error as xray sampler will attempt to fetch rules
xray_sampler = _custom_import_sampler(None, resource=None)
xray_sampler._rules_timer.cancel()
xray_sampler._targets_timer.cancel()
xray_client = xray_sampler._AwsXRayRemoteSampler__xray_client
self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 600)
self.assertEqual(
xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint, "http://localhost:2000/GetSamplingRules"
)

os.environ.pop(OTEL_TRACES_SAMPLER_ARG, None)
os.environ.setdefault(OTEL_TRACES_SAMPLER_ARG, "polling_interval=123")

# May log http request error as xray sampler will attempt to fetch rules
xray_sampler = _custom_import_sampler(None, resource=None)
xray_sampler._rules_timer.cancel()
xray_sampler._targets_timer.cancel()
xray_client = xray_sampler._AwsXRayRemoteSampler__xray_client
self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 123)
self.assertEqual(
xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint, "http://127.0.0.1:2000/GetSamplingRules"
)

os.environ.pop(OTEL_TRACES_SAMPLER_ARG, None)
os.environ.setdefault(OTEL_TRACES_SAMPLER_ARG, "endpoint=https://randomURL:2024")

# May log http request error as xray sampler will attempt to fetch rules
xray_sampler = _custom_import_sampler(None, resource=None)
xray_sampler._rules_timer.cancel()
xray_sampler._targets_timer.cancel()
xray_client = xray_sampler._AwsXRayRemoteSampler__xray_client
self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 300)
self.assertEqual(
xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint, "https://randomURL:2024/GetSamplingRules"
)

def test_import_xray_sampler_with_invalid_environment_arguments(self):
os.environ.pop(OTEL_TRACES_SAMPLER_ARG, None)
os.environ.setdefault(OTEL_TRACES_SAMPLER_ARG, "endpoint=h=tt=p://=loca=lho=st:2000,polling_interval=FOOBAR")

# May log http request error as xray sampler will attempt to fetch rules
xray_sampler = _custom_import_sampler(None, resource=None)
xray_sampler._rules_timer.cancel()
xray_sampler._targets_timer.cancel()
xray_client = xray_sampler._AwsXRayRemoteSampler__xray_client
self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 300)
self.assertEqual(
xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint,
"h=tt=p://=loca=lho=st:2000/GetSamplingRules",
)

os.environ.pop(OTEL_TRACES_SAMPLER_ARG, None)
os.environ.setdefault(OTEL_TRACES_SAMPLER_ARG, ",,=,==,,===,")

# May log http request error as xray sampler will attempt to fetch rules
xray_sampler = _custom_import_sampler(None, resource=None)
xray_sampler._rules_timer.cancel()
xray_sampler._targets_timer.cancel()
xray_client = xray_sampler._AwsXRayRemoteSampler__xray_client
self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 300)
self.assertEqual(
xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint, "http://127.0.0.1:2000/GetSamplingRules"
)

os.environ.pop(OTEL_TRACES_SAMPLER_ARG, None)
os.environ.setdefault(OTEL_TRACES_SAMPLER_ARG, "endpoint,polling_interval")

# May log http request error as xray sampler will attempt to fetch rules
xray_sampler = _custom_import_sampler(None, resource=None)
xray_sampler._rules_timer.cancel()
xray_sampler._targets_timer.cancel()
xray_client = xray_sampler._AwsXRayRemoteSampler__xray_client
self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 300)
self.assertEqual(
xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint, "http://127.0.0.1:2000/GetSamplingRules"
)

0 comments on commit ba62648

Please sign in to comment.