Added credential authentication support for customer sdkstats (#46143)

* Added credential authentication support for customer sdkstats

* Update CHANGELOG

* Address feedback

* Fix test

* Update docstring

Author: rads-1996

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

@@ -7,6 +7,8 @@
 ### Breaking Changes
 
 ### Bugs Fixed
+- Added credential authentication support for customer sdkstats
+  ([#46143](https://github.com/Azure/azure-sdk-for-python/pull/46143))
 
 ### Other Changes
 

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/statsbeat/customer/_customer_sdkstats.py

@@ -19,7 +19,10 @@ def collect_customer_sdkstats(exporter: "BaseExporter") -> None:  # type: ignore
     # Check if already initialized (thread-safe check)
     if not customer_stats.is_initialized:
         # The initialize method is thread-safe and handles double-initialization
-        customer_stats.initialize(connection_string=exporter._connection_string)  # type: ignore
+        customer_stats.initialize(
+            connection_string=exporter._connection_string,  # type: ignore
+            credential=exporter._credential,  # type: ignore
+        )
 
 
 def shutdown_customer_sdkstats_metrics() -> None:

sdk/monitor/azure-monitor-opentelemetry-exporter/azure/monitor/opentelemetry/exporter/statsbeat/customer/_manager.py

@@ -128,11 +128,14 @@ def is_shutdown(self) -> bool:
         """
         return self._status == CustomerSdkStatsStatus.SHUTDOWN  # type: ignore
 
-    def initialize(self, connection_string: str) -> bool:
+    def initialize(self, connection_string: str, credential: Optional[Any] = None) -> bool:
         """Initialize Customer SDKStats collection with the provided connection string.
 
         :param connection_string: Azure Monitor connection string
         :type connection_string: str
+        :param credential: Token credential for AAD authentication. Defaults to None.
+        :type credential: ~azure.core.credentials.TokenCredential or None
+
         :return: True if initialization was successful, False otherwise
         :rtype: bool
         """
@@ -147,24 +150,30 @@ def initialize(self, connection_string: str) -> bool:
                 # Already initialized, return True
                 return True
 
-            return self._do_initialize(connection_string)
+            return self._do_initialize(connection_string, credential=credential)
 
-    def _do_initialize(self, connection_string: str) -> bool:
+    def _do_initialize(self, connection_string: str, credential: Optional[Any] = None) -> bool:
         """Internal initialization method.
 
         :param connection_string: Azure Monitor connection string
         :type connection_string: str
+        :param credential: Token credential for AAD authentication. Defaults to None.
+        :type credential: ~azure.core.credentials.TokenCredential or None
+
         :return: True if initialization was successful, False otherwise
         :rtype: bool
         """
         try:
             # Use delayed import to avoid circular import
             from azure.monitor.opentelemetry.exporter.export.metrics._exporter import AzureMonitorMetricExporter
 
-            self._customer_sdkstats_exporter = AzureMonitorMetricExporter(
-                connection_string=connection_string,
-                is_customer_sdkstats=True,
-            )
+            exporter_kwargs: Dict[str, Any] = {
+                "connection_string": connection_string,
+                "is_customer_sdkstats": True,
+            }
+            if credential is not None:
+                exporter_kwargs["credential"] = credential
+            self._customer_sdkstats_exporter = AzureMonitorMetricExporter(**exporter_kwargs)
             metric_reader_options = {
                 "exporter": self._customer_sdkstats_exporter,
                 "export_interval_millis": get_customer_sdkstats_export_interval() * 1000,  # Default 15m

sdk/monitor/azure-monitor-opentelemetry-exporter/tests/customer_sdk_stats/test_customer_sdkstats.py

@@ -43,6 +43,7 @@ def test_collect_customer_sdkstats(self):
         # Create a mock exporter
         mock_exporter = mock.Mock()
         mock_exporter._connection_string = "InstrumentationKey=12345678-1234-5678-abcd-12345678abcd"
+        mock_exporter._credential = None
 
         # Collect customer SDK stats
         collect_customer_sdkstats(mock_exporter)
@@ -51,6 +52,24 @@ def test_collect_customer_sdkstats(self):
         manager = get_customer_stats_manager()
         self.assertTrue(manager.is_initialized)
 
+    @mock.patch("azure.monitor.opentelemetry.exporter.statsbeat.customer._customer_sdkstats.get_customer_stats_manager")
+    def test_collect_customer_sdkstats_passes_credential(self, mock_get_manager):
+        """Test that credential from exporter is passed to manager.initialize()."""
+        mock_manager = mock.Mock()
+        mock_manager.is_initialized = False
+        mock_get_manager.return_value = mock_manager
+
+        mock_exporter = mock.Mock()
+        mock_exporter._connection_string = "InstrumentationKey=12345678-1234-5678-abcd-12345678abcd"
+        mock_exporter._credential = mock.Mock()
+
+        collect_customer_sdkstats(mock_exporter)
+
+        mock_manager.initialize.assert_called_once_with(
+            connection_string=mock_exporter._connection_string,
+            credential=mock_exporter._credential,
+        )
+
     def test_collect_customer_sdkstats_multiple_calls(self):
         """Test that multiple calls to collect_customer_sdkstats don't cause issues."""
         # Create a mock exporter

sdk/monitor/azure-monitor-opentelemetry-exporter/tests/customer_sdk_stats/test_manager.py

@@ -36,6 +36,7 @@ def setUp(self):
 
         # Get a fresh manager instance
         self.manager = CustomerSdkStatsManager()
+        self.mock_credential = Mock()
 
     def tearDown(self):
         """Clean up test environment."""
@@ -92,16 +93,20 @@ def test_initialize_success(self, mock_meter_provider, mock_metric_reader, mock_
 
         connection_string = "InstrumentationKey=12345678-1234-5678-abcd-12345678abcd"
 
-        # Test initialization
-        result = self.manager.initialize(connection_string)
+        # Test initialization with credential
+        result = self.manager.initialize(connection_string, credential=self.mock_credential)
 
         self.assertTrue(result)
         self.assertEqual(self.manager.status, CustomerSdkStatsStatus.ACTIVE)
         self.assertTrue(self.manager.is_initialized)
         self.assertFalse(self.manager.is_shutdown)
 
         # Verify mocks were called
-        mock_exporter.assert_called_once()
+        mock_exporter.assert_called_once_with(
+            connection_string=connection_string,
+            is_customer_sdkstats=True,
+            credential=self.mock_credential,
+        )
         mock_metric_reader.assert_called_once()
         mock_meter_provider.assert_called_once()
         self.assertEqual(mock_meter.create_observable_gauge.call_count, 3)
@@ -129,6 +134,49 @@ def test_initialize_empty_connection_string(self):
         self.assertFalse(result)
         self.assertEqual(self.manager.status, CustomerSdkStatsStatus.UNINITIALIZED)
 
+    @patch("azure.monitor.opentelemetry.exporter.export.metrics._exporter.AzureMonitorMetricExporter")
+    @patch("azure.monitor.opentelemetry.exporter.statsbeat.customer._manager.PeriodicExportingMetricReader")
+    @patch("azure.monitor.opentelemetry.exporter.statsbeat.customer._manager.MeterProvider")
+    def test_initialize_with_credential(self, mock_meter_provider, mock_metric_reader, mock_exporter):
+        """Test that credential is passed through to the exporter during initialization."""
+        mock_meter = Mock()
+        mock_meter_provider_instance = Mock()
+        mock_meter_provider_instance.get_meter.return_value = mock_meter
+        mock_meter_provider.return_value = mock_meter_provider_instance
+
+        connection_string = "InstrumentationKey=12345678-1234-5678-abcd-12345678abcd"
+        mock_credential = Mock()
+
+        result = self.manager.initialize(connection_string, credential=mock_credential)
+
+        self.assertTrue(result)
+        self.assertTrue(self.manager.is_initialized)
+        mock_exporter.assert_called_once_with(
+            connection_string=connection_string,
+            is_customer_sdkstats=True,
+            credential=mock_credential,
+        )
+
+    @patch("azure.monitor.opentelemetry.exporter.export.metrics._exporter.AzureMonitorMetricExporter")
+    @patch("azure.monitor.opentelemetry.exporter.statsbeat.customer._manager.PeriodicExportingMetricReader")
+    @patch("azure.monitor.opentelemetry.exporter.statsbeat.customer._manager.MeterProvider")
+    def test_initialize_without_credential(self, mock_meter_provider, mock_metric_reader, mock_exporter):
+        """Test that credential is not passed to the exporter when not provided."""
+        mock_meter = Mock()
+        mock_meter_provider_instance = Mock()
+        mock_meter_provider_instance.get_meter.return_value = mock_meter
+        mock_meter_provider.return_value = mock_meter_provider_instance
+
+        connection_string = "InstrumentationKey=12345678-1234-5678-abcd-12345678abcd"
+
+        result = self.manager.initialize(connection_string)
+
+        self.assertTrue(result)
+        mock_exporter.assert_called_once_with(
+            connection_string=connection_string,
+            is_customer_sdkstats=True,
+        )
+
     def test_initialize_multiple_calls(self):
         """Test that multiple initialization calls are handled correctly."""
         with patch("azure.monitor.opentelemetry.exporter.export.metrics._exporter.AzureMonitorMetricExporter"), patch(