stop passing session_id_hex

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

tests/unit/test_client.py

@@ -127,7 +127,6 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
                     connection=connection,
                     execute_response=mock_execute_response,
                     thrift_client=mock_backend,
-                    session_id_hex=Mock(),
                 )
 
                 # Mock execute_command to return our real result set
@@ -187,7 +186,6 @@ def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
             connection=mock_connection,
             execute_response=Mock(),
             thrift_client=mock_backend,
-            session_id_hex=Mock(),
         )
         result_set.results = mock_results
 
@@ -217,7 +215,6 @@ def test_closing_result_set_hard_closes_commands(self):
             mock_connection,
             mock_results_response,
             mock_thrift_backend,
-            session_id_hex=Mock(),
         )
         result_set.results = mock_results
 
@@ -265,9 +262,7 @@ def test_negative_fetch_throws_exception(self):
         mock_backend = Mock()
         mock_backend.fetch_results.return_value = (Mock(), False, 0)
 
-        result_set = ThriftResultSet(
-            Mock(), Mock(), mock_backend
-        )
+        result_set = ThriftResultSet(Mock(), Mock(), mock_backend)
 
         with self.assertRaises(ValueError) as e:
             result_set.fetchmany(-1)

tests/unit/test_fetches.py

@@ -63,7 +63,6 @@ def make_dummy_result_set_from_initial_results(initial_results):
             ),
             thrift_client=mock_thrift_backend,
             t_row_set=None,
-            session_id_hex=Mock(),
         )
         return rs
 
@@ -108,7 +107,6 @@ def fetch_results(
                 is_staging_operation=False,
             ),
             thrift_client=mock_thrift_backend,
-            session_id_hex=Mock(),
         )
         return rs
 

tests/unit/test_telemetry_retry.py

@@ -6,7 +6,8 @@
 from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 from databricks.sql.auth.retry import DatabricksRetryPolicy
 
-PATCH_TARGET = 'urllib3.connectionpool.HTTPSConnectionPool._get_conn'
+PATCH_TARGET = "urllib3.connectionpool.HTTPSConnectionPool._get_conn"
+
 
 def create_mock_conn(responses):
     """Creates a mock connection object whose getresponse() method yields a series of responses."""
@@ -16,15 +17,18 @@ def create_mock_conn(responses):
         mock_http_response = MagicMock()
         mock_http_response.status = resp.get("status")
         mock_http_response.headers = resp.get("headers", {})
-        body = resp.get("body", b'{}')
+        body = resp.get("body", b"{}")
         mock_http_response.fp = io.BytesIO(body)
+
         def release():
             mock_http_response.fp.close()
+
         mock_http_response.release_conn = release
         mock_http_responses.append(mock_http_response)
     mock_conn.getresponse.side_effect = mock_http_responses
     return mock_conn
 
+
 class TestTelemetryClientRetries:
     @pytest.fixture(autouse=True)
     def setup_and_teardown(self):
@@ -49,28 +53,28 @@ def get_client(self, session_id, num_retries=3):
             host_url="test.databricks.com",
         )
         client = TelemetryClientFactory.get_telemetry_client(session_id)
-        
+
         retry_policy = DatabricksRetryPolicy(
             delay_min=0.01,
             delay_max=0.02,
             stop_after_attempts_duration=2.0,
-            stop_after_attempts_count=num_retries, 
+            stop_after_attempts_count=num_retries,
             delay_default=0.1,
             force_dangerous_codes=[],
-            urllib3_kwargs={'total': num_retries}
+            urllib3_kwargs={"total": num_retries},
         )
         adapter = client._http_client.session.adapters.get("https://")
         adapter.max_retries = retry_policy
         return client
 
     @pytest.mark.parametrize(
-    "status_code, description",
-    [
-        (401, "Unauthorized"),
-        (403, "Forbidden"),
-        (501, "Not Implemented"),
-        (200, "Success"),
-    ],
+        "status_code, description",
+        [
+            (401, "Unauthorized"),
+            (403, "Forbidden"),
+            (501, "Not Implemented"),
+            (200, "Success"),
+        ],
     )
     def test_non_retryable_status_codes_are_not_retried(self, status_code, description):
         """
@@ -80,7 +84,9 @@ def test_non_retryable_status_codes_are_not_retried(self, status_code, descripti
         client = self.get_client(f"session-{status_code}")
         mock_responses = [{"status": status_code}]
 
-        with patch(PATCH_TARGET, return_value=create_mock_conn(mock_responses)) as mock_get_conn:
+        with patch(
+            PATCH_TARGET, return_value=create_mock_conn(mock_responses)
+        ) as mock_get_conn:
             client.export_failure_log("TestError", "Test message")
             TelemetryClientFactory.close(client._session_id_hex)
 
@@ -92,16 +98,26 @@ def test_exceeds_retry_count_limit(self):
         Verifies that the client respects the Retry-After header and retries on 429, 502, 503.
         """
         num_retries = 3
-        expected_total_calls = num_retries + 1 
+        expected_total_calls = num_retries + 1
         retry_after = 1
         client = self.get_client("session-exceed-limit", num_retries=num_retries)
-        mock_responses = [{"status": 503, "headers": {"Retry-After": str(retry_after)}}, {"status": 429}, {"status": 502}, {"status": 503}]
-        
-        with patch(PATCH_TARGET, return_value=create_mock_conn(mock_responses)) as mock_get_conn:
+        mock_responses = [
+            {"status": 503, "headers": {"Retry-After": str(retry_after)}},
+            {"status": 429},
+            {"status": 502},
+            {"status": 503},
+        ]
+
+        with patch(
+            PATCH_TARGET, return_value=create_mock_conn(mock_responses)
+        ) as mock_get_conn:
             start_time = time.time()
             client.export_failure_log("TestError", "Test message")
             TelemetryClientFactory.close(client._session_id_hex)
             end_time = time.time()
-            
-            assert mock_get_conn.return_value.getresponse.call_count == expected_total_calls
-            assert end_time - start_time > retry_after
+
+            assert (
+                mock_get_conn.return_value.getresponse.call_count
+                == expected_total_calls
+            )
+            assert end_time - start_time > retry_after