unit tests for initial remote sampler and lint

Author: jj22ee

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/sampler/aws_xray_remote_sampler.py

@@ -21,6 +21,7 @@
 DEFAULT_TARGET_POLLING_INTERVAL = 10
 DEFAULT_SAMPLING_PROXY_ENDPOINT = "http://127.0.0.1:2000"
 
+
 class AwsXRayRemoteSampler(Sampler):
     """
     Remote Sampler for OpenTelemetry that gets sampling configurations from AWS X-Ray
@@ -32,11 +33,17 @@ class AwsXRayRemoteSampler(Sampler):
         log_level: custom log level configuration for remote sampler (Optional)
     """
 
-    __resource : Resource
-    __polling_interval : int
-    __xray_client : AwsXRaySamplingClient
+    __resource: Resource
+    __polling_interval: int
+    __xray_client: AwsXRaySamplingClient
 
-    def __init__(self, resource=None, endpoint=DEFAULT_SAMPLING_PROXY_ENDPOINT, polling_interval=DEFAULT_RULES_POLLING_INTERVAL, log_level = None):
+    def __init__(
+        self,
+        resource=None,
+        endpoint=DEFAULT_SAMPLING_PROXY_ENDPOINT,
+        polling_interval=DEFAULT_RULES_POLLING_INTERVAL,
+        log_level=None,
+    ):
         # Override default log level
         if log_level is not None:
             _logger.setLevel(log_level)
@@ -72,8 +79,7 @@ def __get_and_update_sampling_rules(self):
 
     def __start_sampling_rule_poller(self):
         self.__get_and_update_sampling_rules()
-        # Schedule the next sampling rule poll 
+        # Schedule the next sampling rule poll
         self._timer = Timer(self.__polling_interval, self.__start_sampling_rule_poller)
         self._timer.daemon = True
         self._timer.start()
-

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/sampler/aws_xray_sampling_client.py

@@ -9,6 +9,7 @@
 
 _logger = getLogger(__name__)
 
+
 class AwsXRaySamplingClient:
     def __init__(self, endpoint=None, log_level=None):
         # Override default log level
@@ -21,15 +22,17 @@ def __init__(self, endpoint=None, log_level=None):
 
     def get_sampling_rules(self):
         sampling_rules = []
-        headers = {'content-type': 'application/json'}
+        headers = {"content-type": "application/json"}
 
         try:
             r = requests.post(url=self.__getSamplingRulesEndpoint, headers=headers)
             if r is None:
                 raise Exception("GetSamplingRules response is None")
             sampling_rules_response = r.json()
             if "SamplingRuleRecords" not in sampling_rules_response:
-                raise Exception(f"SamplingRuleRecords is missing in getSamplingRules response:{sampling_rules_response}")
+                raise Exception(
+                    f"SamplingRuleRecords is missing in getSamplingRules response:{sampling_rules_response}"
+                )
 
             sampling_rules_records = sampling_rules_response["SamplingRuleRecords"]
             for record in sampling_rules_records:
@@ -42,4 +45,4 @@ def get_sampling_rules(self):
         except Exception as ex:
             _logger.exception(f"Exception occurred: {ex}")
 
-        return sampling_rules
+        return sampling_rules

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/sampler/sampling_rule.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 class SamplingRule:
     def __init__(
-        self, 
+        self,
         Attributes={},
         FixedRate=None,
         HTTPMethod=None,
@@ -15,18 +15,18 @@ def __init__(
         ServiceName=None,
         ServiceType=None,
         URLPath=None,
-        Version=None
+        Version=None,
     ):
-        self.Attributes=Attributes
-        self.FixedRate=FixedRate
-        self.HTTPMethod=HTTPMethod
-        self.Host=Host
-        self.Priority=Priority
-        self.ReservoirSize=ReservoirSize
-        self.ResourceARN=ResourceARN
-        self.RuleARN=RuleARN
-        self.RuleName=RuleName
-        self.ServiceName=ServiceName
-        self.ServiceType=ServiceType
-        self.URLPath=URLPath
-        self.Version=Version
+        self.Attributes = Attributes
+        self.FixedRate = FixedRate
+        self.HTTPMethod = HTTPMethod
+        self.Host = Host
+        self.Priority = Priority
+        self.ReservoirSize = ReservoirSize
+        self.ResourceARN = ResourceARN
+        self.RuleARN = RuleARN
+        self.RuleName = RuleName
+        self.ServiceName = ServiceName
+        self.ServiceType = ServiceType
+        self.URLPath = URLPath
+        self.Version = Version

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/sampler/data/get-sampling-rules-response-sample.json

@@ -0,0 +1,65 @@
+{
+    "NextToken": null,
+    "SamplingRuleRecords": [
+        {
+            "CreatedAt": 1.67799933E9,
+            "ModifiedAt": 1.67799933E9,
+            "SamplingRule": {
+                "Attributes": {
+                    "foo": "bar",
+                    "doo": "baz"
+                },
+                "FixedRate": 0.05,
+                "HTTPMethod": "*",
+                "Host": "*",
+                "Priority": 1000,
+                "ReservoirSize": 10,
+                "ResourceARN": "*",
+                "RuleARN": "arn:aws:xray:us-west-2:123456789000:sampling-rule/Rule1",
+                "RuleName": "Rule1",
+                "ServiceName": "*",
+                "ServiceType": "AWS::Foo::Bar",
+                "URLPath": "*",
+                "Version": 1
+            }
+        },
+        {
+            "CreatedAt": 0.0,
+            "ModifiedAt": 1.611564245E9,
+            "SamplingRule": {
+                "Attributes": {},
+                "FixedRate": 0.05,
+                "HTTPMethod": "*",
+                "Host": "*",
+                "Priority": 10000,
+                "ReservoirSize": 1,
+                "ResourceARN": "*",
+                "RuleARN": "arn:aws:xray:us-west-2:123456789000:sampling-rule/Default",
+                "RuleName": "Default",
+                "ServiceName": "*",
+                "ServiceType": "*",
+                "URLPath": "*",
+                "Version": 1
+            }
+        },
+        {
+            "CreatedAt": 1.676038494E9,
+            "ModifiedAt": 1.676038494E9,
+            "SamplingRule": {
+                "Attributes": {},
+                "FixedRate": 0.2,
+                "HTTPMethod": "GET",
+                "Host": "*",
+                "Priority": 1,
+                "ReservoirSize": 10,
+                "ResourceARN": "*",
+                "RuleARN": "arn:aws:xray:us-west-2:123456789000:sampling-rule/Rule2",
+                "RuleName": "Rule2",
+                "ServiceName": "FooBar",
+                "ServiceType": "*",
+                "URLPath": "/foo/bar",
+                "Version": 1
+            }
+        }
+    ]
+}

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/sampler/test_aws_xray_remote_sampler.py

@@ -0,0 +1,45 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from logging import DEBUG
+from unittest import TestCase
+
+from amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler import AwsXRayRemoteSampler
+from opentelemetry.sdk.resources import Resource
+
+
+class AwsXRayRemoteSamplerTest(TestCase):
+    def test_create_remote_sampler_with_empty_resource(self):
+        rs = AwsXRayRemoteSampler(resource=Resource.get_empty())
+        self.assertIsNotNone(rs._timer)
+        self.assertEqual(rs._AwsXRayRemoteSampler__polling_interval, 300)
+        self.assertIsNotNone(rs._AwsXRayRemoteSampler__xray_client)
+        self.assertIsNotNone(rs._AwsXRayRemoteSampler__resource)
+
+    def test_create_remote_sampler_with_populated_resource(self):
+        rs = AwsXRayRemoteSampler(
+            resource=Resource.create({"service.name": "test-service-name", "cloud.platform": "test-cloud-platform"})
+        )
+        self.assertIsNotNone(rs._timer)
+        self.assertEqual(rs._AwsXRayRemoteSampler__polling_interval, 300)
+        self.assertIsNotNone(rs._AwsXRayRemoteSampler__xray_client)
+        self.assertIsNotNone(rs._AwsXRayRemoteSampler__resource)
+        self.assertEqual(rs._AwsXRayRemoteSampler__resource.attributes["service.name"], "test-service-name")
+        self.assertEqual(rs._AwsXRayRemoteSampler__resource.attributes["cloud.platform"], "test-cloud-platform")
+
+    def test_create_remote_sampler_with_all_fields_populated(self):
+        rs = AwsXRayRemoteSampler(
+            resource=Resource.create({"service.name": "test-service-name", "cloud.platform": "test-cloud-platform"}),
+            endpoint="http://abc.com",
+            polling_interval=120,
+            log_level=DEBUG,
+        )
+        self.assertIsNotNone(rs._timer)
+        self.assertEqual(rs._AwsXRayRemoteSampler__polling_interval, 120)
+        self.assertIsNotNone(rs._AwsXRayRemoteSampler__xray_client)
+        self.assertIsNotNone(rs._AwsXRayRemoteSampler__resource)
+        self.assertEqual(
+            rs._AwsXRayRemoteSampler__xray_client._AwsXRaySamplingClient__getSamplingRulesEndpoint,
+            "http://abc.com/GetSamplingRules",
+        )  # "http://127.0.0.1:2000"
+        self.assertEqual(rs._AwsXRayRemoteSampler__resource.attributes["service.name"], "test-service-name")
+        self.assertEqual(rs._AwsXRayRemoteSampler__resource.attributes["cloud.platform"], "test-cloud-platform")

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/sampler/test_aws_xray_sampling_client.py

@@ -0,0 +1,42 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import json
+import os
+from logging import getLogger
+from unittest import TestCase
+from unittest.mock import patch
+
+from amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler import AwsXRayRemoteSampler
+from amazon.opentelemetry.distro.sampler.aws_xray_sampling_client import AwsXRaySamplingClient
+
+SAMPLING_CLIENT_LOGGER_NAME = "amazon.opentelemetry.distro.sampler.aws_xray_sampling_client"
+_logger = getLogger(SAMPLING_CLIENT_LOGGER_NAME)
+
+TEST_DIR = os.path.dirname(os.path.realpath(__file__))
+DATA_DIR = os.path.join(TEST_DIR, "data")
+
+
+class AwsXRaySamplingClientTest(TestCase):
+    @patch("requests.post")
+    def test_get_no_sampling_rules(self, mock_post=None, prom_rw=None):
+        mock_post.return_value.configure_mock(**{"json.return_value": {"SamplingRuleRecords": []}})
+        client = AwsXRaySamplingClient("http://127.0.0.1:2000")
+        sampling_rules = client.get_sampling_rules()
+        self.assertTrue(len(sampling_rules) == 0)
+
+    @patch("requests.post")
+    def test_get_invalid_response(self, mock_post=None, prom_rw=None):
+        mock_post.return_value.configure_mock(**{"json.return_value": {}})
+        client = AwsXRaySamplingClient("http://127.0.0.1:2000")
+        with self.assertLogs(_logger, level="ERROR") as cm:
+            sampling_rules = client.get_sampling_rules()
+        self.assertTrue(len(sampling_rules) == 0)
+
+    @patch("requests.post")
+    def test_get_two_sampling_rules(self, mock_post=None, prom_rw=None):
+        with open(f"{DATA_DIR}/get-sampling-rules-response-sample.json") as f:
+            mock_post.return_value.configure_mock(**{"json.return_value": json.load(f)})
+            f.close()
+        client = AwsXRaySamplingClient("http://127.0.0.1:2000")
+        sampling_rules = client.get_sampling_rules()
+        self.assertTrue(len(sampling_rules) == 3)