test

Signed-off-by: Sai Shree Pradhan <[email protected]>

Author: saishreeeee

src/databricks/sql/client.py

@@ -270,7 +270,10 @@ def read(self) -> Optional[OAuthToken]:
                 error_message=str(e),
                 host_url=server_hostname,
                 http_path=http_path,
-                port=self.session.port,
+                port=kwargs.get("_port", 443),
+                user_agent=self.session.useragent_header
+                if hasattr(self, "session")
+                else None,
             )
             raise e
 

src/databricks/sql/telemetry/telemetry_client.py

@@ -469,3 +469,4 @@ def connection_failure_log(
         telemetry_client._user_agent = user_agent
 
         telemetry_client.export_failure_log(error_name, error_message)
+        TelemetryClientFactory.close(UNAUTH_DUMMY_SESSION_ID)

tests/unit/test_telemetry.py

@@ -8,7 +8,6 @@
     NoopTelemetryClient,
     TelemetryClientFactory,
     TelemetryHelper,
-    BaseTelemetryClient,
 )
 from databricks.sql.telemetry.models.enums import AuthMech, AuthFlow
 from databricks.sql.auth.authenticators import (
@@ -290,3 +289,37 @@ def test_factory_shutdown_flow(self):
         TelemetryClientFactory.close(session2)
         assert TelemetryClientFactory._initialized is False
         assert TelemetryClientFactory._executor is None
+
+    @patch("databricks.sql.client.Session")
+    @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
+    def test_connection_failure_sends_correct_telemetry_payload(
+        self, mock_send_telemetry, mock_session
+    ):
+        """
+        Verify that a connection failure constructs and sends the correct
+        telemetry payload via _send_telemetry.
+        """
+
+        error_message = "Could not connect to host"
+        mock_session.side_effect = Exception(error_message)
+
+        try:
+            from databricks.sql.client import Connection
+            Connection(server_hostname="test-host", http_path="/test-path")
+        except Exception as e:
+            assert str(e) == error_message
+
+        mock_send_telemetry.assert_called_once()
+        
+        call_arguments = mock_send_telemetry.call_args.args
+        sent_events = call_arguments[0]
+        
+        assert len(sent_events) == 1
+        telemetry_log = sent_events[0]
+        
+        assert telemetry_log.entry.sql_driver_log is not None
+        
+        error_info = telemetry_log.entry.sql_driver_log.error_info
+        assert error_info is not None
+        assert error_info.error_name == "Exception"
+        assert error_info.stack_trace == error_message