Add tracer_provider kwarg to AzureMonitorTraceExporter (Azure#36363)

Author: avramdj

sdk/monitor/azure-monitor-opentelemetry-exporter/CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ### Features Added
 
+- Allow passing in of custom `TracerProvider` for `AzureMonitorTraceExporter`
+    ([#36363](https://github.com/Azure/azure-sdk-for-python/pull/36363))
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/export/trace/_exporter.py

@@ -71,6 +71,10 @@
 class AzureMonitorTraceExporter(BaseExporter, SpanExporter):
     """Azure Monitor Trace exporter for OpenTelemetry."""
 
+    def __init__(self, **kwargs: Any):
+        self._tracer_provider = kwargs.pop("tracer_provider", get_tracer_provider())
+        super().__init__(**kwargs)
+
     def export(self, spans: Sequence[ReadableSpan], **kwargs: Any) -> SpanExportResult: # pylint: disable=unused-argument
         """Export span data.
 
@@ -83,7 +87,7 @@ def export(self, spans: Sequence[ReadableSpan], **kwargs: Any) -> SpanExportResu
         if spans and self._should_collect_otel_resource_metric():
             resource = None
             try:
-                tracer_provider = get_tracer_provider()
+                tracer_provider = self._tracer_provider
                 resource = tracer_provider.resource # type: ignore
                 envelopes.append(self._get_otel_resource_envelope(resource))
             except AttributeError as e:

sdk/monitor/azure-monitor-opentelemetry-exporter/tests/trace/test_trace.py

@@ -10,6 +10,7 @@
 from unittest import mock
 
 # pylint: disable=import-error
+from opentelemetry.trace import get_tracer_provider, set_tracer_provider
 from opentelemetry.sdk import trace, resources
 from opentelemetry.sdk.trace.export import SpanExportResult
 from opentelemetry.sdk.util.instrumentation import InstrumentationScope
@@ -55,17 +56,43 @@ def setUpClass(cls):
         ] = "1234abcd-5678-4efa-8abc-1234567890ab"
         os.environ["APPLICATIONINSIGHTS_STATSBEAT_DISABLED_ALL"] = "true"
         os.environ["APPLICATIONINSIGHTS_OPENTELEMETRY_RESOURCE_METRIC_DISABLED"] = "true"
-        cls._exporter = AzureMonitorTraceExporter()
+        cls._exporter = AzureMonitorTraceExporter(
+            tracer_provider=trace.TracerProvider(),
+        )
 
     @classmethod
     def tearDownClass(cls):
         shutil.rmtree(cls._exporter.storage._path, True)
 
     def test_constructor(self):
         """Test the constructor."""
+        tp = trace.TracerProvider()
+        set_tracer_provider(tp)
         exporter = AzureMonitorTraceExporter(
             connection_string="InstrumentationKey=4321abcd-5678-4efa-8abc-1234567890ab",
         )
+        self.assertEqual(
+            exporter._tracer_provider,
+            tp,
+        )
+        self.assertEqual(
+            exporter._instrumentation_key,
+            "4321abcd-5678-4efa-8abc-1234567890ab",
+        )
+
+    def test_constructor_tracer_provider(self):
+        """Test the constructor."""
+        tp = trace.TracerProvider()
+        tp2 = trace.TracerProvider()
+        set_tracer_provider(tp)
+        exporter = AzureMonitorTraceExporter(
+            connection_string="InstrumentationKey=4321abcd-5678-4efa-8abc-1234567890ab",
+            tracer_provider=tp2,
+        )
+        self.assertEqual(
+            exporter._tracer_provider,
+            tp2,
+        )
         self.assertEqual(
             exporter._instrumentation_key,
             "4321abcd-5678-4efa-8abc-1234567890ab",
@@ -1326,7 +1353,6 @@ def test_export_otel_resource_metric(self, mock_get_tracer_provider):
         del os.environ["APPLICATIONINSIGHTS_OPENTELEMETRY_RESOURCE_METRIC_DISABLED"]
         mock_tracer_provider = mock.Mock()
         mock_get_tracer_provider.return_value = mock_tracer_provider
-        exporter = self._exporter
         test_resource = resources.Resource(
             attributes={
                 "string_test_key": "string_value",
@@ -1347,6 +1373,9 @@ def test_export_otel_resource_metric(self, mock_get_tracer_provider):
         )
         test_span.start()
         test_span.end()
+        exporter = AzureMonitorTraceExporter(
+            _disable_offline_storage=True,
+        )
         with mock.patch(
             "azure.monitor.opentelemetry.exporter.AzureMonitorTraceExporter._transmit"
         ) as transmit:  # noqa: E501
@@ -1438,4 +1467,4 @@ def test_check_instrumentation_span_azure_sdk(self):
             "azure.monitor.opentelemetry.exporter._utils.add_instrumentation"
         ) as add:
             _check_instrumentation_span(span)
-            add.assert_called_once_with(_AZURE_SDK_OPENTELEMETRY_NAME)
+            add.assert_called_once_with(_AZURE_SDK_OPENTELEMETRY_NAME)