tests

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

tests/unit/test_telemetry.py

@@ -7,14 +7,14 @@
     NoopTelemetryClient,
     TelemetryClientFactory,
     TelemetryHelper,
-    BaseTelemetryClient,
 )
 from databricks.sql.telemetry.models.enums import AuthMech, AuthFlow
 from databricks.sql.auth.authenticators import (
     AccessTokenAuthProvider,
     DatabricksOAuthProvider,
     ExternalAuthProvider,
 )
+from databricks import sql
 
 
 @pytest.fixture
@@ -311,8 +311,6 @@ def test_connection_failure_sends_correct_telemetry_payload(
         mock_session.side_effect = Exception(error_message)
 
         try:
-            from databricks import sql
-
             sql.connect(server_hostname="test-host", http_path="/test-path")
         except Exception as e:
             assert str(e) == error_message
@@ -321,3 +319,90 @@ def test_connection_failure_sends_correct_telemetry_payload(
         call_arguments = mock_export_failure_log.call_args
         assert call_arguments[0][0] == "Exception"
         assert call_arguments[0][1] == error_message
+
+
+@patch("databricks.sql.client.Session")
+class TestTelemetryFeatureFlag:
+    """Tests the interaction between the telemetry feature flag and connection parameters."""
+
+    def _mock_ff_response(self, mock_requests_get, enabled: bool):
+        """Helper to configure the mock response for the feature flag endpoint."""
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        payload = {
+            "flags": [
+                {
+                    "name": "databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForPythonDriver",
+                    "value": str(enabled).lower(),
+                }
+            ],
+            "ttl_seconds": 3600,
+        }
+        mock_response.json.return_value = payload
+        mock_requests_get.return_value = mock_response
+
+    @patch("databricks.sql.common.feature_flag.requests.get")
+    def test_telemetry_enabled_when_flag_is_true(
+        self, mock_requests_get, MockSession
+    ):
+        """Telemetry should be ON when enable_telemetry=True and server flag is 'true'."""
+        self._mock_ff_response(mock_requests_get, enabled=True)
+        mock_session_instance = MockSession.return_value
+        mock_session_instance.guid_hex = "test-session-ff-true"
+        mock_session_instance.auth_provider = AccessTokenAuthProvider("token")
+
+        conn = sql.client.Connection(
+            server_hostname="test",
+            http_path="test",
+            access_token="test",
+            enable_telemetry=True,
+        )
+
+        assert conn.telemetry_enabled is True
+        mock_requests_get.assert_called_once()
+        client = TelemetryClientFactory.get_telemetry_client("test-session-ff-true")
+        assert isinstance(client, TelemetryClient)
+
+    @patch("databricks.sql.common.feature_flag.requests.get")
+    def test_telemetry_disabled_when_flag_is_false(
+        self, mock_requests_get, MockSession
+    ):
+        """Telemetry should be OFF when enable_telemetry=True but server flag is 'false'."""
+        self._mock_ff_response(mock_requests_get, enabled=False)
+        mock_session_instance = MockSession.return_value
+        mock_session_instance.guid_hex = "test-session-ff-false"
+        mock_session_instance.auth_provider = AccessTokenAuthProvider("token")
+
+        conn = sql.client.Connection(
+            server_hostname="test",
+            http_path="test",
+            access_token="test",
+            enable_telemetry=True,
+        )
+
+        assert conn.telemetry_enabled is False
+        mock_requests_get.assert_called_once()
+        client = TelemetryClientFactory.get_telemetry_client("test-session-ff-false")
+        assert isinstance(client, NoopTelemetryClient)
+
+    @patch("databricks.sql.common.feature_flag.requests.get")
+    def test_telemetry_disabled_when_flag_request_fails(
+        self, mock_requests_get, MockSession
+    ):
+        """Telemetry should default to OFF if the feature flag network request fails."""
+        mock_requests_get.side_effect = Exception("Network is down")
+        mock_session_instance = MockSession.return_value
+        mock_session_instance.guid_hex = "test-session-ff-fail"
+        mock_session_instance.auth_provider = AccessTokenAuthProvider("token")
+
+        conn = sql.client.Connection(
+            server_hostname="test",
+            http_path="test",
+            access_token="test",
+            enable_telemetry=True,
+        )
+
+        assert conn.telemetry_enabled is False
+        mock_requests_get.assert_called_once()
+        client = TelemetryClientFactory.get_telemetry_client("test-session-ff-fail")
+        assert isinstance(client, NoopTelemetryClient)