Skip to content

Commit 4ba5ae2

Browse files
srprashsrprash
committed
adding unit tests
1 parent a8a895c commit 4ba5ae2

File tree

5 files changed

+169
-41
lines changed

5 files changed

+169
-41
lines changed

‎aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,8 @@
1616
AwsMetricAttributesSpanExporterBuilder,
1717
)
1818
from amazon.opentelemetry.distro.aws_span_metrics_processor_builder import AwsSpanMetricsProcessorBuilder
19+
from amazon.opentelemetry.distro.otlp_udp_exporter import OtlpUdpMetricExporter, OtlpUdpSpanExporter
1920
from amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler import AwsXRayRemoteSampler
20-
from amazon.opentelemetry.distro.otlp_udp_exporter import OtlpUdpMetricExporter
21-
from amazon.opentelemetry.distro.otlp_udp_exporter import OtlpUdpSpanExporter
2221
from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter as OTLPHttpOTLPMetricExporter
2322
from opentelemetry.sdk._configuration import (
2423
_get_exporter_names,
@@ -64,6 +63,7 @@
6463
APPLICATION_SIGNALS_EXPORTER_ENDPOINT_CONFIG = "OTEL_AWS_APPLICATION_SIGNALS_EXPORTER_ENDPOINT"
6564
METRIC_EXPORT_INTERVAL_CONFIG = "OTEL_METRIC_EXPORT_INTERVAL"
6665
DEFAULT_METRIC_EXPORT_INTERVAL = 60000.0
66+
AWS_LAMBDA_FUNCTION_NAME_ENV_VAR = "AWS_LAMBDA_FUNCTION_NAME"
6767

6868
_logger: Logger = getLogger(__name__)
6969

@@ -268,7 +268,7 @@ def _is_application_signals_enabled():
268268

269269
def _is_lambda_environment():
270270
# detect if running in AWS Lambda environment
271-
return "AWS_LAMBDA_FUNCTION_NAME" in os.environ
271+
return AWS_LAMBDA_FUNCTION_NAME_ENV_VAR in os.environ
272272

273273

274274
class ApplicationSignalsExporterProvider:
@@ -300,7 +300,7 @@ def create_exporter(self):
300300

301301
if _is_lambda_environment():
302302
# When running in Lambda, export Application Signals metrics over UDP
303-
return OtlpUdpMetricExporter(preferred_temporality=temporality_dict)
303+
return OtlpUdpMetricExporter(endpoint="127.0.0.1:2000", preferred_temporality=temporality_dict)
304304

305305
if protocol == "http/protobuf":
306306
application_signals_endpoint = os.environ.get(

‎aws-opentelemetry-distro/src/amazon/opentelemetry/distro/otlp_udp_exporter.py

Lines changed: 34 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,67 +1,66 @@
1-
from typing import Optional, Sequence, Dict
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
import socket
4+
from typing import Dict, Optional, Sequence
5+
6+
from typing_extensions import override
27

8+
from opentelemetry.exporter.otlp.proto.common.metrics_encoder import encode_metrics
9+
from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans
310
from opentelemetry.sdk.metrics._internal.aggregation import AggregationTemporality
11+
from opentelemetry.sdk.metrics._internal.export import MetricExportResult
12+
from opentelemetry.sdk.metrics._internal.point import MetricsData
13+
from opentelemetry.sdk.metrics.export import MetricExporter
414
from opentelemetry.sdk.metrics.view import Aggregation
515
from opentelemetry.sdk.trace import ReadableSpan
616
from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
7-
from opentelemetry.sdk.metrics.export import MetricExporter
8-
from typing_extensions import override
9-
from opentelemetry.sdk.metrics._internal.point import MetricsData
10-
from opentelemetry.sdk.metrics._internal.export import MetricExportResult
11-
from opentelemetry.exporter.otlp.proto.common.metrics_encoder import (
12-
encode_metrics,
13-
)
14-
from opentelemetry.exporter.otlp.proto.common.trace_encoder import (
15-
encode_spans,
16-
)
1717

18-
import socket
18+
from logging import getLogger, Logger
1919

2020
DEFAULT_ENDPOINT = "127.0.0.1:2000"
2121
PROTOCOL_HEADER = '{"format":"json","version":1}\n'
22-
PROTOCOL_DELIMITER = '\n'
22+
PROTOCOL_DELIMITER = "\n"
2323

24+
_logger: Logger = getLogger(__name__)
2425

2526
class UdpExporter:
26-
def __init__(
27-
self,
28-
endpoint: Optional[str] = None
29-
):
30-
self._endpoint = endpoint or DEFAULT_ENDPOINT
27+
def __init__(self, endpoint: Optional[str] = None):
28+
self._endpoint = endpoint or DEFAULT_ENDPOINT # TODO: read from some env var??
3129
self._host, self._port = self._parse_endpoint(self._endpoint)
3230
self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
3331
self._socket.setblocking(False)
3432

35-
def send_data(self, data: str, format: str):
36-
udp_data = f'{{"format":"{format}","data":{data}}}'
33+
def send_data(self, data: str, signal_format: str):
34+
udp_data = f'{{"format":"{signal_format}","data":{data}}}'
3735
message = PROTOCOL_HEADER + udp_data
3836

3937
try:
4038
print("Sending UDP data: ", message) # TODO: remove
41-
self._socket.sendto(message.encode('utf-8'), (self._host, int(self._port)))
42-
except Exception as e:
43-
print("Error sending UDP data: ", e)
39+
self._socket.sendto(message.encode("utf-8"), (self._host, int(self._port)))
40+
except Exception as exc: # pylint: disable=broad-except
41+
_logger.error("Error sending UDP data: %s", exc)
4442

4543
def shutdown(self):
4644
self._socket.close()
4745

46+
# pylint: disable=no-self-use
4847
def _parse_endpoint(self, endpoint: str) -> tuple[str, int]:
4948
try:
5049
vals = endpoint.split(":")
5150
host = vals[0]
5251
port = int(vals[1])
53-
except Exception as e:
54-
raise ValueError(f"Invalid endpoint: {endpoint}") from e
52+
except Exception as exc: # pylint: disable=broad-except
53+
raise ValueError(f"Invalid endpoint: {endpoint}") from exc
5554

5655
return host, port
5756

5857

5958
class OtlpUdpMetricExporter(MetricExporter):
6059
def __init__(
61-
self,
62-
endpoint: Optional[str] = None,
63-
preferred_temporality: Dict[type, AggregationTemporality] = None,
64-
preferred_aggregation: Dict[type, Aggregation] = None
60+
self,
61+
endpoint: Optional[str] = None,
62+
preferred_temporality: Dict[type, AggregationTemporality] = None,
63+
preferred_aggregation: Dict[type, Aggregation] = None,
6564
):
6665
super().__init__(
6766
preferred_temporality=preferred_temporality,
@@ -71,13 +70,13 @@ def __init__(
7170

7271
@override
7372
def export(
74-
self,
75-
metrics_data: MetricsData,
76-
timeout_millis: float = 10_000,
77-
**kwargs,
73+
self,
74+
metrics_data: MetricsData,
75+
timeout_millis: float = 10_000,
76+
**kwargs,
7877
) -> MetricExportResult:
7978
serialized_data = encode_metrics(metrics_data).SerializeToString()
80-
self._udp_exporter.send_data(data=serialized_data, format="OTEL_V1_METRICS") # TODO: Convert to constant
79+
self._udp_exporter.send_data(data=serialized_data, signal_format="OTEL_V1_METRICS") # TODO: Convert to constant
8180
return MetricExportResult.SUCCESS # TODO: send appropriate status back. Need to??
8281

8382
def force_flush(self, timeout_millis: float = 10_000) -> bool:
@@ -94,7 +93,7 @@ def __init__(self, endpoint: Optional[str] = None):
9493
@override
9594
def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
9695
serialized_data = encode_spans(spans).SerializeToString()
97-
self._udp_exporter.send_data(data=serialized_data, format="OTEL_V1_TRACES") # TODO: Convert to constant
96+
self._udp_exporter.send_data(data=serialized_data, signal_format="OTEL_V1_TRACES") # TODO: Convert to constant
9897
return SpanExportResult.SUCCESS # TODO: send appropriate status back. Need to??
9998

10099
@override

‎aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,15 @@
1919
)
2020
from amazon.opentelemetry.distro.aws_opentelemetry_distro import AwsOpenTelemetryDistro
2121
from amazon.opentelemetry.distro.aws_span_metrics_processor import AwsSpanMetricsProcessor
22+
from amazon.opentelemetry.distro.otlp_udp_exporter import OtlpUdpSpanExporter
2223
from amazon.opentelemetry.distro.sampler._aws_xray_sampling_client import _AwsXRaySamplingClient
2324
from amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler import AwsXRayRemoteSampler
2425
from opentelemetry.environment_variables import OTEL_LOGS_EXPORTER, OTEL_METRICS_EXPORTER, OTEL_TRACES_EXPORTER
2526
from opentelemetry.exporter.otlp.proto.common._internal.metrics_encoder import OTLPMetricExporterMixin
2627
from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter as OTLPGrpcOTLPMetricExporter
2728
from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter as OTLPHttpOTLPMetricExporter
2829
from opentelemetry.sdk.environment_variables import OTEL_TRACES_SAMPLER, OTEL_TRACES_SAMPLER_ARG
30+
from opentelemetry.sdk.metrics._internal.export import MetricExporter
2931
from opentelemetry.sdk.resources import Resource
3032
from opentelemetry.sdk.trace import Span, SpanProcessor, Tracer, TracerProvider
3133
from opentelemetry.sdk.trace.export import SpanExporter
@@ -254,6 +256,16 @@ def test_customize_exporter(self):
254256
self.assertEqual(mock_exporter, customized_exporter._delegate)
255257
os.environ.pop("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", None)
256258

259+
# when Application Signals is enabled and running in lambda
260+
os.environ.setdefault("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", "True")
261+
os.environ.setdefault("AWS_LAMBDA_FUNCTION_NAME", "myLambdaFunc")
262+
customized_exporter = _customize_exporter(mock_exporter, Resource.get_empty())
263+
self.assertNotEqual(mock_exporter, customized_exporter)
264+
self.assertIsInstance(customized_exporter, AwsMetricAttributesSpanExporter)
265+
self.assertIsInstance(customized_exporter._delegate, OtlpUdpSpanExporter)
266+
os.environ.pop("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", None)
267+
os.environ.pop("AWS_LAMBDA_FUNCTION_NAME", None)
268+
257269
def test_customize_span_processors(self):
258270
mock_tracer_provider: TracerProvider = MagicMock()
259271
_customize_span_processors(mock_tracer_provider, Resource.get_empty())
@@ -286,6 +298,12 @@ def test_application_signals_exporter_provider(self):
286298
self.assertIsInstance(exporter, OTLPHttpOTLPMetricExporter)
287299
self.assertEqual("http://localhost:4316/v1/metrics", exporter._endpoint)
288300

301+
# When in Lambda, exporter should be UDP.
302+
os.environ.setdefault("AWS_LAMBDA_FUNCTION_NAME", "myLambdaFunc")
303+
exporter: MetricExporter = ApplicationSignalsExporterProvider().create_exporter()
304+
self.assertIsInstance(exporter, OtlpUdpSpanExporter)
305+
self.assertEqual("127.0.0.1:2000", exporter._udp_exporter._endpoint)
306+
289307

290308
def validate_distro_environ():
291309
tc: TestCase = TestCase()

‎aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_otlp_udp_exporter.py

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
import socket
4+
import unittest
5+
from unittest import TestCase
6+
from unittest.mock import MagicMock, patch
7+
8+
from amazon.opentelemetry.distro.otlp_udp_exporter import (
9+
DEFAULT_ENDPOINT,
10+
PROTOCOL_HEADER,
11+
OtlpUdpMetricExporter,
12+
OtlpUdpSpanExporter,
13+
UdpExporter,
14+
)
15+
from opentelemetry.sdk.metrics._internal.export import MetricExportResult
16+
from opentelemetry.sdk.trace.export import SpanExportResult
17+
18+
19+
class TestUdpExporter(TestCase):
20+
21+
@patch("amazon.opentelemetry.distro.otlp_udp_exporter.socket.socket")
22+
def test_udp_exporter_init_default(self, mock_socket):
23+
exporter = UdpExporter()
24+
self.assertEqual(exporter._endpoint, DEFAULT_ENDPOINT)
25+
self.assertEqual(exporter._host, "127.0.0.1")
26+
self.assertEqual(exporter._port, 2000)
27+
mock_socket.assert_called_once_with(socket.AF_INET, socket.SOCK_DGRAM)
28+
mock_socket().setblocking.assert_called_once_with(False)
29+
30+
@patch("amazon.opentelemetry.distro.otlp_udp_exporter.socket.socket")
31+
def test_udp_exporter_init_with_endpoint(self, mock_socket):
32+
exporter = UdpExporter(endpoint="localhost:5000")
33+
self.assertNotEqual(exporter._endpoint, DEFAULT_ENDPOINT)
34+
self.assertEqual(exporter._host, "localhost")
35+
self.assertEqual(exporter._port, 5000)
36+
mock_socket.assert_called_once_with(socket.AF_INET, socket.SOCK_DGRAM)
37+
mock_socket().setblocking.assert_called_once_with(False)
38+
39+
@patch("amazon.opentelemetry.distro.otlp_udp_exporter.socket.socket")
40+
def test_udp_exporter_init_invalid_endpoint(self, mock_socket):
41+
with self.assertRaises(ValueError):
42+
UdpExporter(endpoint="invalidEndpoint:port")
43+
44+
# pylint: disable=no-self-use
45+
@patch("amazon.opentelemetry.distro.otlp_udp_exporter.socket.socket")
46+
def test_send_data(self, mock_socket):
47+
mock_socket_instance = mock_socket.return_value
48+
exporter = UdpExporter()
49+
exporter.send_data('encoded_data', "signal")
50+
expected_message = PROTOCOL_HEADER + '{"format":"signal","data":encoded_data}'
51+
mock_socket_instance.sendto.assert_called_once_with(expected_message.encode("utf-8"), ("127.0.0.1", 2000))
52+
53+
@patch("amazon.opentelemetry.distro.otlp_udp_exporter.socket.socket")
54+
def test_shutdown(self, mock_socket):
55+
mock_socket_instance = mock_socket.return_value
56+
exporter = UdpExporter()
57+
exporter.shutdown()
58+
mock_socket_instance.close.assert_called_once()
59+
60+
61+
class TestOtlpUdpMetricExporter(unittest.TestCase):
62+
63+
@patch("amazon.opentelemetry.distro.otlp_udp_exporter.encode_metrics")
64+
@patch("amazon.opentelemetry.distro.otlp_udp_exporter.UdpExporter")
65+
def test_export(self, mock_udp_exporter, mock_encode_metrics):
66+
mock_udp_exporter_instance = mock_udp_exporter.return_value
67+
mock_encoded_data = MagicMock()
68+
mock_encode_metrics.return_value.SerializeToString.return_value = mock_encoded_data
69+
exporter = OtlpUdpMetricExporter()
70+
result = exporter.export(MagicMock())
71+
mock_udp_exporter_instance.send_data.assert_called_once_with(
72+
data=mock_encoded_data, signal_format="OTEL_V1_METRICS"
73+
)
74+
self.assertEqual(result, MetricExportResult.SUCCESS)
75+
76+
# pylint: disable=no-self-use
77+
@patch("amazon.opentelemetry.distro.otlp_udp_exporter.UdpExporter")
78+
def test_shutdown(self, mock_udp_exporter):
79+
mock_udp_exporter_instance = mock_udp_exporter.return_value
80+
exporter = OtlpUdpMetricExporter()
81+
exporter.shutdown()
82+
mock_udp_exporter_instance.shutdown.assert_called_once()
83+
84+
85+
class TestOtlpUdpSpanExporter(unittest.TestCase):
86+
87+
@patch("amazon.opentelemetry.distro.otlp_udp_exporter.encode_spans")
88+
@patch("amazon.opentelemetry.distro.otlp_udp_exporter.UdpExporter")
89+
def test_export(self, mock_udp_exporter, mock_encode_spans):
90+
mock_udp_exporter_instance = mock_udp_exporter.return_value
91+
mock_encoded_data = MagicMock()
92+
mock_encode_spans.return_value.SerializeToString.return_value = mock_encoded_data
93+
exporter = OtlpUdpSpanExporter()
94+
result = exporter.export(MagicMock())
95+
mock_udp_exporter_instance.send_data.assert_called_once_with(
96+
data=mock_encoded_data, signal_format="OTEL_V1_TRACES"
97+
)
98+
self.assertEqual(result, SpanExportResult.SUCCESS)
99+
100+
# pylint: disable=no-self-use
101+
@patch("amazon.opentelemetry.distro.otlp_udp_exporter.UdpExporter")
102+
def test_shutdown(self, mock_udp_exporter):
103+
mock_udp_exporter_instance = mock_udp_exporter.return_value
104+
exporter = OtlpUdpSpanExporter()
105+
exporter.shutdown()
106+
mock_udp_exporter_instance.shutdown.assert_called_once()
107+
108+
109+
# TODO: remove this line for final PR
110+
if __name__ == "__main__":
111+
unittest.main()

‎test_lambda.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
# test script for appsignals lambda behavior
22

3-
from opentelemetry import trace
4-
53
from time import sleep
64

5+
from opentelemetry import trace
6+
77
if __name__ == "__main__":
88
tracer = trace.get_tracer("test-tracer")
99
with tracer.start_as_current_span("parent"):

0 commit comments

Comments
 (0)