fix: Update tests to use host_url instead of session_id_hex

- Update circuit breaker test to check logger.debug instead of logger.info
- Replace all session_id_hex test parameters with host_url
- Apply Black formatting to exc.py and telemetry_client.py

This fixes test failures caused by the signature change from session_id_hex
to host_url in the Error class and TelemetryClientFactory.

Author: samikshya-db

src/databricks/sql/exc.py

@@ -11,9 +11,7 @@ class Error(Exception):
     `context`: Optional extra context about the error. MUST be JSON serializable
     """
 
-    def __init__(
-        self, message=None, context=None, host_url=None, *args, **kwargs
-    ):
+    def __init__(self, message=None, context=None, host_url=None, *args, **kwargs):
         super().__init__(message, *args, **kwargs)
         self.message = message
         self.context = context or {}

src/databricks/sql/telemetry/telemetry_client.py

@@ -586,13 +586,13 @@ def initialize_telemetry_client(
                             batch_size=batch_size,
                             client_context=client_context,
                         )
-                        TelemetryClientFactory._clients[host_url] = _TelemetryClientHolder(
-                            client
-                        )
+                        TelemetryClientFactory._clients[
+                            host_url
+                        ] = _TelemetryClientHolder(client)
                     else:
-                        TelemetryClientFactory._clients[host_url] = _TelemetryClientHolder(
-                            NoopTelemetryClient()
-                        )
+                        TelemetryClientFactory._clients[
+                            host_url
+                        ] = _TelemetryClientHolder(NoopTelemetryClient())
         except Exception as e:
             logger.debug("Failed to initialize telemetry client: %s", e)
             # Fallback to NoopTelemetryClient to ensure connection doesn't fail
@@ -648,7 +648,8 @@ def close(host_url):
             # Only close if no more references
             if remaining_refs <= 0:
                 logger.debug(
-                    "Closing telemetry client for host %s (no more references)", host_url
+                    "Closing telemetry client for host %s (no more references)",
+                    host_url,
                 )
                 TelemetryClientFactory._clients.pop(host_url, None)
                 holder.client.close()

tests/unit/test_circuit_breaker_manager.py

@@ -157,4 +157,4 @@ def test_circuit_breaker_state_listener_transitions(self, old_state, new_state):
 
         with patch("databricks.sql.telemetry.circuit_breaker_manager.logger") as mock_logger:
             listener.state_change(mock_cb, mock_old_state, mock_new_state)
-            mock_logger.info.assert_called()
+            mock_logger.debug.assert_called()

tests/unit/test_client.py

@@ -714,7 +714,7 @@ def test_autocommit_setter_wraps_database_error(self, mock_session_class):
         server_error = DatabaseError(
             "AUTOCOMMIT_SET_DURING_ACTIVE_TRANSACTION",
             context={"sql_state": "25000"},
-            session_id_hex="test-session-id",
+            host_url="test-session-id",
         )
         mock_cursor.execute.side_effect = server_error
 
@@ -737,7 +737,7 @@ def test_autocommit_setter_preserves_exception_chain(self, mock_session_class):
 
         mock_cursor = Mock()
         original_error = DatabaseError(
-            "Original error", session_id_hex="test-session-id"
+            "Original error", host_url="test-session-id"
         )
         mock_cursor.execute.side_effect = original_error
 
@@ -772,7 +772,7 @@ def test_commit_wraps_database_error(self, mock_session_class):
         server_error = DatabaseError(
             "MULTI_STATEMENT_TRANSACTION_NO_ACTIVE_TRANSACTION",
             context={"sql_state": "25000"},
-            session_id_hex="test-session-id",
+            host_url="test-session-id",
         )
         mock_cursor.execute.side_effect = server_error
 
@@ -822,7 +822,7 @@ def test_rollback_wraps_database_error(self, mock_session_class):
         server_error = DatabaseError(
             "Unexpected rollback error",
             context={"sql_state": "HY000"},
-            session_id_hex="test-session-id",
+            host_url="test-session-id",
         )
         mock_cursor.execute.side_effect = server_error
 

tests/unit/test_download_manager.py

@@ -20,7 +20,7 @@ def create_download_manager(
             max_download_threads,
             lz4_compressed,
             ssl_options=SSLOptions(),
-            session_id_hex=Mock(),
+            host_url=Mock(),
             statement_id=Mock(),
             chunk_id=0,
             http_client=mock_http_client,

tests/unit/test_downloader.py

@@ -57,7 +57,7 @@ def test_run_link_expired(self, mock_time):
             result_link,
             ssl_options=SSLOptions(),
             chunk_id=0,
-            session_id_hex=Mock(),
+            host_url=Mock(),
             statement_id=Mock(),
             http_client=mock_http_client,
         )
@@ -80,7 +80,7 @@ def test_run_link_past_expiry_buffer(self, mock_time):
             result_link,
             ssl_options=SSLOptions(),
             chunk_id=0,
-            session_id_hex=Mock(),
+            host_url=Mock(),
             statement_id=Mock(),
             http_client=mock_http_client,
         )
@@ -107,7 +107,7 @@ def test_run_get_response_not_ok(self, mock_time):
             result_link,
             ssl_options=SSLOptions(),
             chunk_id=0,
-            session_id_hex=Mock(),
+            host_url=Mock(),
             statement_id=Mock(),
             http_client=mock_http_client,
         )
@@ -137,7 +137,7 @@ def test_run_uncompressed_successful(self, mock_time):
                 result_link,
                 ssl_options=SSLOptions(),
                 chunk_id=0,
-                session_id_hex=Mock(),
+                host_url=Mock(),
                 statement_id=Mock(),
                 http_client=mock_http_client,
             )
@@ -170,7 +170,7 @@ def test_run_compressed_successful(self, mock_time):
                 result_link,
                 ssl_options=SSLOptions(),
                 chunk_id=0,
-                session_id_hex=Mock(),
+                host_url=Mock(),
                 statement_id=Mock(),
                 http_client=mock_http_client,
             )
@@ -194,7 +194,7 @@ def test_download_connection_error(self, mock_time):
             result_link,
             ssl_options=SSLOptions(),
             chunk_id=0,
-            session_id_hex=Mock(),
+            host_url=Mock(),
             statement_id=Mock(),
             http_client=mock_http_client,
         )
@@ -216,7 +216,7 @@ def test_download_timeout(self, mock_time):
             result_link,
             ssl_options=SSLOptions(),
             chunk_id=0,
-            session_id_hex=Mock(),
+            host_url=Mock(),
             statement_id=Mock(),
             http_client=mock_http_client,
         )

tests/unit/test_telemetry.py

@@ -46,7 +46,7 @@ def mock_telemetry_client():
     ):
         return TelemetryClient(
             telemetry_enabled=True,
-            session_id_hex=session_id,
+            host_url=session_id,
             auth_provider=auth_provider,
             host_url="test-host.com",
             executor=executor,
@@ -242,7 +242,7 @@ def test_client_lifecycle_flow(self):
         ):
             TelemetryClientFactory.initialize_telemetry_client(
                 telemetry_enabled=True,
-                session_id_hex=session_id_hex,
+                host_url=session_id_hex,
                 auth_provider=auth_provider,
                 host_url="test-host.com",
                 batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
@@ -267,7 +267,7 @@ def test_disabled_telemetry_creates_noop_client(self):
 
         TelemetryClientFactory.initialize_telemetry_client(
             telemetry_enabled=False,
-            session_id_hex=session_id_hex,
+            host_url=session_id_hex,
             auth_provider=None,
             host_url="test-host.com",
             batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
@@ -289,7 +289,7 @@ def test_factory_error_handling(self):
         ):
             TelemetryClientFactory.initialize_telemetry_client(
                 telemetry_enabled=True,
-                session_id_hex=session_id,
+                host_url=session_id,
                 auth_provider=AccessTokenAuthProvider("token"),
                 host_url="test-host.com",
                 batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
@@ -313,7 +313,7 @@ def test_factory_shutdown_flow(self):
             for session in [session1, session2]:
                 TelemetryClientFactory.initialize_telemetry_client(
                     telemetry_enabled=True,
-                    session_id_hex=session,
+                    host_url=session,
                     auth_provider=AccessTokenAuthProvider("token"),
                     host_url="test-host.com",
                     batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,