fix imports in tests

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

Author: varun-edachali-dbx

tests/unit/test_client.py

@@ -26,7 +26,7 @@
 from databricks.sql.types import Row
 from databricks.sql.result_set import ResultSet, ThriftResultSet
 from databricks.sql.backend.types import CommandId, CommandState
-from databricks.sql.utils import ExecuteResponse
+from databricks.sql.backend.types import ExecuteResponse
 
 from tests.unit.test_fetches import FetchTests
 from tests.unit.test_thrift_backend import ThriftBackendTestSuite

tests/unit/test_fetches.py

@@ -8,7 +8,8 @@
     pa = None
 
 import databricks.sql.client as client
-from databricks.sql.utils import ExecuteResponse, ArrowQueue
+from databricks.sql.backend.types import ExecuteResponse
+from databricks.sql.utils import ArrowQueue
 from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
 from databricks.sql.result_set import ThriftResultSet
 
@@ -48,7 +49,7 @@ def make_dummy_result_set_from_initial_results(initial_results):
                 description=Mock(),
                 lz4_compressed=Mock(),
                 command_id=None,
-                arrow_queue=arrow_queue,
+                results_queue=arrow_queue,
                 arrow_schema_bytes=schema.serialize().to_pybytes(),
                 is_staging_operation=False,
             ),
@@ -97,7 +98,7 @@ def fetch_results(
                 ],
                 lz4_compressed=Mock(),
                 command_id=None,
-                arrow_queue=None,
+                results_queue=None,
                 arrow_schema_bytes=None,
                 is_staging_operation=False,
             ),

tests/unit/test_fetches_bench.py

@@ -10,7 +10,8 @@
 import pytest
 
 import databricks.sql.client as client
-from databricks.sql.utils import ExecuteResponse, ArrowQueue
+from databricks.sql.backend.types import ExecuteResponse
+from databricks.sql.utils import ArrowQueue
 
 
 @pytest.mark.skipif(pa is None, reason="PyArrow is not installed")

tests/unit/test_sea_backend.py

@@ -10,7 +10,7 @@
 from unittest.mock import patch, MagicMock, Mock
 
 from databricks.sql.backend.sea.backend import SeaDatabricksClient
-from databricks.sql.backend.sea_result_set import SeaResultSet
+from databricks.sql.result_set import SeaResultSet
 from databricks.sql.backend.types import SessionId, CommandId, CommandState, BackendType
 from databricks.sql.types import SSLOptions
 from databricks.sql.auth.authenticators import AuthProvider