Fix all tests

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

src/databricks/sql/client.py

@@ -443,7 +443,7 @@ def get_protocol_version(openSessionResp: TOpenSessionResp):
     @property
     def open(self) -> bool:
         """Return whether the connection is open by checking if the session is open."""
-        return self.session.is_open
+        return hasattr(self, 'session') and self.session.is_open
 
     def cursor(
         self,

tests/unit/test_auth.py

@@ -294,7 +294,7 @@ def test_get_token_failure(self, token_source, http_response):
             mock_http_client.execute.return_value.__enter__.return_value = mock_response
             mock_http_client.execute.return_value.__exit__.return_value = None
 
-            with pytest.raises(Exception):
+            with pytest.raises(Exception) as e:
                 token_source.get_token()
             assert "Failed to get token: 400" in str(e.value)
 

tests/unit/test_sea_queue.py

@@ -7,7 +7,7 @@
 """
 
 import pytest
-from unittest.mock import Mock, patch
+from unittest.mock import Mock, patch, MagicMock
 
 from databricks.sql.backend.sea.queue import (
     JsonQueue,
@@ -184,6 +184,7 @@ def description(self):
     def test_build_queue_json_array(self, json_manifest, sample_data):
         """Test building a JSON array queue."""
         result_data = ResultData(data=sample_data)
+        mock_http_client = MagicMock()
 
         queue = SeaResultSetQueueFactory.build_queue(
             result_data=result_data,
@@ -194,6 +195,7 @@ def test_build_queue_json_array(self, json_manifest, sample_data):
             max_download_threads=10,
             sea_client=Mock(),
             lz4_compressed=False,
+            http_client=mock_http_client,
         )
 
         assert isinstance(queue, JsonQueue)
@@ -217,6 +219,8 @@ def test_build_queue_arrow_stream(
         ]
         result_data = ResultData(data=None, external_links=external_links)
 
+        mock_http_client = MagicMock()
+
         with patch(
             "databricks.sql.backend.sea.queue.ResultFileDownloadManager"
         ), patch.object(SeaCloudFetchQueue, "_create_next_table", return_value=None):
@@ -229,13 +233,15 @@ def test_build_queue_arrow_stream(
                 max_download_threads=10,
                 sea_client=mock_sea_client,
                 lz4_compressed=False,
+                http_client=mock_http_client,
             )
 
         assert isinstance(queue, SeaCloudFetchQueue)
 
     def test_build_queue_invalid_format(self, invalid_manifest):
         """Test building a queue with invalid format."""
         result_data = ResultData(data=[])
+        mock_http_client = MagicMock()
 
         with pytest.raises(ProgrammingError, match="Invalid result format"):
             SeaResultSetQueueFactory.build_queue(
@@ -247,6 +253,7 @@ def test_build_queue_invalid_format(self, invalid_manifest):
                 max_download_threads=10,
                 sea_client=Mock(),
                 lz4_compressed=False,
+                http_client=mock_http_client,
             )
 
 
@@ -339,6 +346,7 @@ def test_init_with_valid_initial_link(
     ):
         """Test initialization with valid initial link."""
         # Create a queue with valid initial link
+        mock_http_client = MagicMock()
         with patch.object(SeaCloudFetchQueue, "_create_next_table", return_value=None):
             queue = SeaCloudFetchQueue(
                 result_data=ResultData(external_links=[sample_external_link]),
@@ -349,6 +357,7 @@ def test_init_with_valid_initial_link(
                 total_chunk_count=1,
                 lz4_compressed=False,
                 description=description,
+                http_client=mock_http_client,
             )
 
         # Verify attributes
@@ -367,6 +376,7 @@ def test_init_no_initial_links(
     ):
         """Test initialization with no initial links."""
         # Create a queue with empty initial links
+        mock_http_client = MagicMock()
         queue = SeaCloudFetchQueue(
             result_data=ResultData(external_links=[]),
             max_download_threads=5,
@@ -376,6 +386,7 @@ def test_init_no_initial_links(
             total_chunk_count=0,
             lz4_compressed=False,
             description=description,
+            http_client=mock_http_client,
         )
         assert queue.table is None
 
@@ -462,7 +473,7 @@ def test_hybrid_disposition_with_attachment(
         # Create result data with attachment
         attachment_data = b"mock_arrow_data"
         result_data = ResultData(attachment=attachment_data)
-
+        mock_http_client = MagicMock()
         # Build queue
         queue = SeaResultSetQueueFactory.build_queue(
             result_data=result_data,
@@ -473,6 +484,7 @@ def test_hybrid_disposition_with_attachment(
             max_download_threads=10,
             sea_client=mock_sea_client,
             lz4_compressed=False,
+            http_client=mock_http_client,
         )
 
         # Verify ArrowQueue was created
@@ -508,7 +520,8 @@ def test_hybrid_disposition_with_external_links(
         # Create result data with external links but no attachment
         result_data = ResultData(external_links=external_links, attachment=None)
 
-        # Build queue
+        # Build queue   
+        mock_http_client = MagicMock()
         queue = SeaResultSetQueueFactory.build_queue(
             result_data=result_data,
             manifest=arrow_manifest,
@@ -518,6 +531,7 @@ def test_hybrid_disposition_with_external_links(
             max_download_threads=10,
             sea_client=mock_sea_client,
             lz4_compressed=False,
+            http_client=mock_http_client,
         )
 
         # Verify SeaCloudFetchQueue was created
@@ -548,7 +562,7 @@ def test_hybrid_disposition_with_compressed_attachment(
 
         # Create result data with attachment
         result_data = ResultData(attachment=compressed_data)
-
+        mock_http_client = MagicMock()
         # Build queue with lz4_compressed=True
         queue = SeaResultSetQueueFactory.build_queue(
             result_data=result_data,
@@ -559,6 +573,7 @@ def test_hybrid_disposition_with_compressed_attachment(
             max_download_threads=10,
             sea_client=mock_sea_client,
             lz4_compressed=True,
+            http_client=mock_http_client,
         )
 
         # Verify ArrowQueue was created with decompressed data

tests/unit/test_session.py

@@ -75,8 +75,9 @@ def test_http_header_passthrough(self, mock_client_class):
         call_kwargs = mock_client_class.call_args[1]
         assert ("foo", "bar") in call_kwargs["http_headers"]
 
+    @patch("%s.client.UnifiedHttpClient" % PACKAGE_NAME)
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
-    def test_tls_arg_passthrough(self, mock_client_class):
+    def test_tls_arg_passthrough(self, mock_client_class, mock_http_client):
         databricks.sql.connect(
             **self.DUMMY_CONNECTION_ARGS,
             _tls_verify_hostname="hostname",