mock out pollers

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

@@ -3,13 +3,15 @@
 import os
 import time
 from unittest import TestCase
+from unittest.mock import patch
 
 from amazon.opentelemetry.distro.aws_opentelemetry_configurator import (
     AwsOpenTelemetryConfigurator,
     _custom_import_sampler,
 )
 from amazon.opentelemetry.distro.aws_opentelemetry_distro import AwsOpenTelemetryDistro
 from amazon.opentelemetry.distro.sampler._aws_xray_sampling_client import _AwsXRaySamplingClient
+from amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler import AwsXRayRemoteSampler
 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
@@ -58,30 +60,30 @@ def test_trace_id_ratio_sampler(self):
             # 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
+    # Test method for importing xray sampler
     # Cannot test this logic via `aws_otel_configurator.configure()` because that will
-    # attempt to setup tracer provider again, which may cause issue
+    # attempt to setup tracer provider again, which can be only be done once (already done)
+    @patch.object(AwsXRayRemoteSampler, "_AwsXRayRemoteSampler__start_sampling_rule_poller", lambda x: None)
+    @patch.object(AwsXRayRemoteSampler, "_AwsXRayRemoteSampler__start_sampling_target_poller", lambda x: None)
     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: Sampler = _custom_import_sampler(None, resource=None)
-        xray_sampler._rules_timer.cancel()
-        xray_sampler._targets_timer.cancel()
         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"
         )
 
+    @patch.object(AwsXRayRemoteSampler, "_AwsXRayRemoteSampler__start_sampling_rule_poller", lambda x: None)
+    @patch.object(AwsXRayRemoteSampler, "_AwsXRayRemoteSampler__start_sampling_target_poller", lambda x: None)
     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: Sampler = _custom_import_sampler(None, resource=None)
-        xray_sampler._rules_timer.cancel()
-        xray_sampler._targets_timer.cancel()
         xray_client: _AwsXRaySamplingClient = xray_sampler._AwsXRayRemoteSampler__xray_client
         self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 600)
         self.assertEqual(
@@ -93,35 +95,32 @@ def test_import_xray_sampler_with_valid_environment_arguments(self):
 
         # May log http request error as xray sampler will attempt to fetch rules
         xray_sampler: Sampler = _custom_import_sampler(None, resource=None)
-        xray_sampler._rules_timer.cancel()
-        xray_sampler._targets_timer.cancel()
         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"
         )
 
         os.environ.pop(OTEL_TRACES_SAMPLER_ARG, None)
-        os.environ.setdefault(OTEL_TRACES_SAMPLER_ARG, "endpoint=https://randomURL:2024")
+        os.environ.setdefault(OTEL_TRACES_SAMPLER_ARG, "endpoint=http://cloudwatch-agent.amazon-cloudwatch:2000")
 
         # May log http request error as xray sampler will attempt to fetch rules
         xray_sampler: Sampler = _custom_import_sampler(None, resource=None)
-        xray_sampler._rules_timer.cancel()
-        xray_sampler._targets_timer.cancel()
         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"
+            xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint,
+            "http://cloudwatch-agent.amazon-cloudwatch:2000/GetSamplingRules",
         )
 
+    @patch.object(AwsXRayRemoteSampler, "_AwsXRayRemoteSampler__start_sampling_rule_poller", lambda x: None)
+    @patch.object(AwsXRayRemoteSampler, "_AwsXRayRemoteSampler__start_sampling_target_poller", lambda x: None)
     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: Sampler = _custom_import_sampler(None, resource=None)
-        xray_sampler._rules_timer.cancel()
-        xray_sampler._targets_timer.cancel()
         xray_client: _AwsXRaySamplingClient = xray_sampler._AwsXRayRemoteSampler__xray_client
         self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 300)
         self.assertEqual(
@@ -134,8 +133,6 @@ def test_import_xray_sampler_with_invalid_environment_arguments(self):
 
         # May log http request error as xray sampler will attempt to fetch rules
         xray_sampler: Sampler = _custom_import_sampler(None, resource=None)
-        xray_sampler._rules_timer.cancel()
-        xray_sampler._targets_timer.cancel()
         xray_client: _AwsXRaySamplingClient = xray_sampler._AwsXRayRemoteSampler__xray_client
         self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 300)
         self.assertEqual(
@@ -147,8 +144,6 @@ def test_import_xray_sampler_with_invalid_environment_arguments(self):
 
         # May log http request error as xray sampler will attempt to fetch rules
         xray_sampler: Sampler = _custom_import_sampler(None, resource=None)
-        xray_sampler._rules_timer.cancel()
-        xray_sampler._targets_timer.cancel()
         xray_client: _AwsXRaySamplingClient = xray_sampler._AwsXRayRemoteSampler__xray_client
         self.assertEqual(xray_sampler._AwsXRayRemoteSampler__polling_interval, 300)
         self.assertEqual(