Skip to content

Commit

Permalink
typing
Browse files Browse the repository at this point in the history
  • Loading branch information
jj22ee committed Feb 16, 2024
1 parent ba62648 commit bb7fd73
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 18 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -143,15 +143,15 @@ def _init_tracing(
set_tracer_provider(trace_provider)


def _custom_import_sampler(sampler_name, resource) -> Sampler:
def _custom_import_sampler(sampler_name: str, resource: 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
sampler_argument_env: str = os.getenv(OTEL_TRACES_SAMPLER_ARG, None)
endpoint: str = None
polling_interval: int = None

if sampler_argument_env is not None:
args = sampler_argument_env.split(",")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,11 @@
_custom_import_sampler,
)
from amazon.opentelemetry.distro.aws_opentelemetry_distro import AwsOpenTelemetryDistro
from amazon.opentelemetry.distro.sampler._aws_xray_sampling_client import _AwsXRaySamplingClient
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
from opentelemetry.sdk.trace import Span, Tracer, TracerProvider
from opentelemetry.sdk.trace.sampling import Sampler
from opentelemetry.trace import get_tracer_provider


Expand Down Expand Up @@ -63,10 +65,10 @@ 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: Sampler = _custom_import_sampler(None, resource=None)
xray_sampler._rules_timer.cancel()
xray_sampler._targets_timer.cancel()
xray_client = xray_sampler._AwsXRayRemoteSampler__xray_client
xray_client: _AwsXRaySamplingClient = 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"
Expand All @@ -77,10 +79,10 @@ def test_import_xray_sampler_with_valid_environment_arguments(self):
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: Sampler = _custom_import_sampler(None, resource=None)
xray_sampler._rules_timer.cancel()
xray_sampler._targets_timer.cancel()
xray_client = xray_sampler._AwsXRayRemoteSampler__xray_client
xray_client: _AwsXRaySamplingClient = 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"
Expand All @@ -90,10 +92,10 @@ def test_import_xray_sampler_with_valid_environment_arguments(self):
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: Sampler = _custom_import_sampler(None, resource=None)
xray_sampler._rules_timer.cancel()
xray_sampler._targets_timer.cancel()
xray_client = xray_sampler._AwsXRayRemoteSampler__xray_client
xray_client: _AwsXRaySamplingClient = 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"
Expand All @@ -103,10 +105,10 @@ def test_import_xray_sampler_with_valid_environment_arguments(self):
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: Sampler = _custom_import_sampler(None, resource=None)
xray_sampler._rules_timer.cancel()
xray_sampler._targets_timer.cancel()
xray_client = xray_sampler._AwsXRayRemoteSampler__xray_client
xray_client: _AwsXRaySamplingClient = 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"
Expand All @@ -117,10 +119,10 @@ def test_import_xray_sampler_with_invalid_environment_arguments(self):
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: Sampler = _custom_import_sampler(None, resource=None)
xray_sampler._rules_timer.cancel()
xray_sampler._targets_timer.cancel()
xray_client = xray_sampler._AwsXRayRemoteSampler__xray_client
xray_client: _AwsXRaySamplingClient = xray_sampler._AwsXRayRemoteSampler__xray_client
self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 300)
self.assertEqual(
xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint,
Expand All @@ -131,10 +133,10 @@ def test_import_xray_sampler_with_invalid_environment_arguments(self):
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: Sampler = _custom_import_sampler(None, resource=None)
xray_sampler._rules_timer.cancel()
xray_sampler._targets_timer.cancel()
xray_client = xray_sampler._AwsXRayRemoteSampler__xray_client
xray_client: _AwsXRaySamplingClient = 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"
Expand All @@ -144,10 +146,10 @@ def test_import_xray_sampler_with_invalid_environment_arguments(self):
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: Sampler = _custom_import_sampler(None, resource=None)
xray_sampler._rules_timer.cancel()
xray_sampler._targets_timer.cancel()
xray_client = xray_sampler._AwsXRayRemoteSampler__xray_client
xray_client: _AwsXRaySamplingClient = 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"
Expand Down

0 comments on commit bb7fd73

Please sign in to comment.