fix: telemetry appinsights connection string provider

packages/uipath/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath"
-version = "2.10.9"
+version = "2.10.10"
 description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools."
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

packages/uipath/src/uipath/telemetry/__init__.py

@@ -1,8 +1,17 @@
 from ._track import (  # noqa: D104
     flush_events,
     is_telemetry_enabled,
+    reset_event_client,
+    set_event_connection_string_provider,
     track,
     track_event,
 )
 
-__all__ = ["track", "track_event", "is_telemetry_enabled", "flush_events"]
+__all__ = [
+    "track",
+    "track_event",
+    "is_telemetry_enabled",
+    "flush_events",
+    "set_event_connection_string_provider",
+    "reset_event_client",
+]

packages/uipath/src/uipath/telemetry/_track.py

@@ -4,7 +4,7 @@
 from functools import wraps
 from importlib.metadata import version
 from logging import INFO, WARNING, LogRecord, getLogger
-from typing import Any, Callable, Dict, Mapping, Optional, Union
+from typing import Any, Callable, ClassVar, Dict, Mapping, Optional, Union
 
 from opentelemetry.sdk._logs import LoggingHandler
 from opentelemetry.util.types import AnyValue
@@ -217,6 +217,18 @@ class _AppInsightsEventClient:
     _initialized = False
     _client: Optional[Any] = None
     _atexit_registered = False
+    _connection_string_provider: ClassVar[Optional[Callable[[], Optional[str]]]] = None
+
+    @staticmethod
+    def set_connection_string_provider(
+        provider: Callable[[], Optional[str]],
+    ) -> None:
+        """Override how the connection string is resolved.
+
+        Args:
+            provider: Zero-arg callable returning a connection string or None.
+        """
+        _AppInsightsEventClient._connection_string_provider = provider
 
     @staticmethod
     def _initialize() -> None:
@@ -234,7 +246,10 @@ def _initialize() -> None:
         if not _HAS_APPINSIGHTS:
             return
 
-        connection_string = _get_connection_string()
+        if _AppInsightsEventClient._connection_string_provider:
+            connection_string = _AppInsightsEventClient._connection_string_provider()
+        else:
+            connection_string = _get_connection_string()
         if not connection_string:
             return
 
@@ -330,6 +345,13 @@ def register_atexit_flush() -> None:
             atexit.register(_AppInsightsEventClient.flush)
             _AppInsightsEventClient._atexit_registered = True
 
+    @staticmethod
+    def reset() -> None:
+        """Flush pending events and reset so the next call re-initializes."""
+        _AppInsightsEventClient.flush()
+        _AppInsightsEventClient._client = None
+        _AppInsightsEventClient._initialized = False
+
 
 class _TelemetryClient:
     """A class to handle telemetry using OpenTelemetry for method tracking."""
@@ -456,6 +478,22 @@ def flush_events() -> None:
     _AppInsightsEventClient.flush()
 
 
+def set_event_connection_string_provider(
+    provider: Callable[[], Optional[str]],
+) -> None:
+    """Override how the Application Insights connection string is resolved.
+
+    Args:
+        provider: Zero-arg callable returning a connection string or None.
+    """
+    _AppInsightsEventClient.set_connection_string_provider(provider)
+
+
+def reset_event_client() -> None:
+    """Flush pending events and reset so the next ``track_event`` re-initializes."""
+    _AppInsightsEventClient.reset()
+
+
 def track_cli_event(
     name: str,
     properties: Optional[Dict[str, Any]] = None,

packages/uipath/tests/telemetry/test_track.py

@@ -9,6 +9,8 @@
     _TelemetryClient,
     flush_events,
     is_telemetry_enabled,
+    reset_event_client,
+    set_event_connection_string_provider,
     track,
     track_event,
 )
@@ -81,11 +83,13 @@ def setup_method(self):
         """Reset AppInsightsEventClient state before each test."""
         _AppInsightsEventClient._initialized = False
         _AppInsightsEventClient._client = None
+        _AppInsightsEventClient._connection_string_provider = None
 
     def teardown_method(self):
         """Clean up after each test."""
         _AppInsightsEventClient._initialized = False
         _AppInsightsEventClient._client = None
+        _AppInsightsEventClient._connection_string_provider = None
 
     @patch("uipath.telemetry._track._CONNECTION_STRING", "$CONNECTION_STRING")
     def test_initialize_no_connection_string(self):
@@ -255,6 +259,164 @@ def test_flush_no_client(self):
         # Should not raise any exception
         _AppInsightsEventClient.flush()
 
+    @patch("uipath.telemetry._track.TelemetryChannel")
+    @patch("uipath.telemetry._track.SynchronousQueue")
+    @patch("uipath.telemetry._track._DiagnosticSender")
+    @patch("uipath.telemetry._track._HAS_APPINSIGHTS", True)
+    @patch("uipath.telemetry._track.AppInsightsTelemetryClient")
+    def test_connection_string_provider_overrides_default(
+        self, mock_client_class, mock_sender_class, mock_queue_class, mock_channel_class
+    ):
+        """Test that a custom provider is used instead of _get_connection_string."""
+        mock_client = MagicMock()
+        mock_client_class.return_value = mock_client
+
+        def provider() -> str:
+            return (
+                "InstrumentationKey=from-provider;IngestionEndpoint=https://custom.com/"
+            )
+
+        _AppInsightsEventClient.set_connection_string_provider(provider)
+
+        _AppInsightsEventClient._initialize()
+
+        assert _AppInsightsEventClient._client is mock_client
+        mock_client_class.assert_called_once_with(
+            "from-provider", telemetry_channel=mock_channel_class.return_value
+        )
+        mock_sender_class.assert_called_once_with(
+            service_endpoint_uri="https://custom.com/v2/track"
+        )
+
+    @patch("uipath.telemetry._track._HAS_APPINSIGHTS", True)
+    def test_connection_string_provider_returning_none_skips_client(self):
+        """Test that provider returning None results in no client."""
+        _AppInsightsEventClient.set_connection_string_provider(lambda: None)
+
+        _AppInsightsEventClient._initialize()
+
+        assert _AppInsightsEventClient._initialized is True
+        assert _AppInsightsEventClient._client is None
+
+    @patch("uipath.telemetry._track.TelemetryChannel")
+    @patch("uipath.telemetry._track.SynchronousQueue")
+    @patch("uipath.telemetry._track._DiagnosticSender")
+    @patch("uipath.telemetry._track._HAS_APPINSIGHTS", True)
+    @patch("uipath.telemetry._track.AppInsightsTelemetryClient")
+    @patch(
+        "uipath.telemetry._track._CONNECTION_STRING",
+        "InstrumentationKey=builtin-key",
+    )
+    def test_provider_bypasses_builtin_fallback(
+        self, mock_client_class, mock_sender_class, mock_queue_class, mock_channel_class
+    ):
+        """Test that provider prevents fallback to _CONNECTION_STRING."""
+        mock_client = MagicMock()
+        mock_client_class.return_value = mock_client
+
+        _AppInsightsEventClient.set_connection_string_provider(
+            lambda: "InstrumentationKey=provider-key"
+        )
+
+        with patch.dict(os.environ, {}, clear=True):
+            _AppInsightsEventClient._initialize()
+
+        # Should use provider-key, not builtin-key
+        mock_client_class.assert_called_once_with(
+            "provider-key", telemetry_channel=mock_channel_class.return_value
+        )
+
+    def test_reset_clears_initialized_and_client(self):
+        """Test that reset clears initialized flag and client."""
+        _AppInsightsEventClient._initialized = True
+        _AppInsightsEventClient._client = MagicMock()
+
+        _AppInsightsEventClient.reset()
+
+        assert _AppInsightsEventClient._initialized is False
+        assert _AppInsightsEventClient._client is None
+
+    def test_reset_flushes_before_clearing(self):
+        """Test that reset flushes pending events before clearing."""
+        mock_client = MagicMock()
+        _AppInsightsEventClient._initialized = True
+        _AppInsightsEventClient._client = mock_client
+
+        _AppInsightsEventClient.reset()
+
+        mock_client.flush.assert_called_once()
+
+    @patch("uipath.telemetry._track.TelemetryChannel")
+    @patch("uipath.telemetry._track.SynchronousQueue")
+    @patch("uipath.telemetry._track._DiagnosticSender")
+    @patch("uipath.telemetry._track._HAS_APPINSIGHTS", True)
+    @patch("uipath.telemetry._track.AppInsightsTelemetryClient")
+    def test_reset_allows_reinitialization_with_new_connection_string(
+        self, mock_client_class, mock_sender_class, mock_queue_class, mock_channel_class
+    ):
+        """Test that after reset, next initialize reads current env."""
+        mock_client_1 = MagicMock()
+        mock_client_2 = MagicMock()
+        mock_client_class.side_effect = [mock_client_1, mock_client_2]
+
+        # First init with connection string A
+        with patch.dict(
+            os.environ,
+            {"TELEMETRY_CONNECTION_STRING": "InstrumentationKey=key-a"},
+        ):
+            _AppInsightsEventClient._initialize()
+
+        assert _AppInsightsEventClient._client is mock_client_1
+
+        # Reset
+        _AppInsightsEventClient.reset()
+
+        # Second init with connection string B
+        with patch.dict(
+            os.environ,
+            {"TELEMETRY_CONNECTION_STRING": "InstrumentationKey=key-b"},
+        ):
+            _AppInsightsEventClient._initialize()
+
+        assert _AppInsightsEventClient._client is mock_client_2
+        assert mock_client_class.call_count == 2
+
+
+class TestPublicProviderAndResetFunctions:
+    """Test the public set_event_connection_string_provider and reset_event_client."""
+
+    def setup_method(self) -> None:
+        """Reset state before each test."""
+        _AppInsightsEventClient._initialized = False
+        _AppInsightsEventClient._client = None
+        _AppInsightsEventClient._connection_string_provider = None
+
+    def teardown_method(self) -> None:
+        """Clean up after each test."""
+        _AppInsightsEventClient._initialized = False
+        _AppInsightsEventClient._client = None
+        _AppInsightsEventClient._connection_string_provider = None
+
+    def test_set_event_connection_string_provider_sets_provider(self) -> None:
+        """Test that the public function sets the provider on the client."""
+
+        def provider() -> str:
+            return "InstrumentationKey=test"
+
+        set_event_connection_string_provider(provider)
+
+        assert _AppInsightsEventClient._connection_string_provider is provider
+
+    def test_reset_event_client_resets_state(self) -> None:
+        """Test that the public function resets client state."""
+        _AppInsightsEventClient._initialized = True
+        _AppInsightsEventClient._client = MagicMock()
+
+        reset_event_client()
+
+        assert _AppInsightsEventClient._initialized is False
+        assert _AppInsightsEventClient._client is None
+
 
 class TestTelemetryClient:
     """Test _TelemetryClient functionality."""