Refractor

databricks_sql_connector_core/tests/unit/test_fetches_bench.py

@@ -1,15 +1,20 @@
 import unittest
 from unittest.mock import Mock
 
-import pyarrow as pa
 import uuid
 import time
 import pytest
 
 import databricks.sql.client as client
 from databricks.sql.utils import ExecuteResponse, ArrowQueue
+from tests.e2e.predicate import pysql_supports_arrow
 
+try:
+    import pyarrow as pa
+except ImportError:
+    pa = None
 
+@pytest.mark.skipif(not pysql_supports_arrow(), reason="Skipping because pyarrow is not installed")
 class FetchBenchmarkTests(unittest.TestCase):
     """
     Micro benchmark test for Arrow result handling.

databricks_sql_connector_core/tests/unit/test_thrift_backend.py

@@ -2,18 +2,22 @@
 from decimal import Decimal
 import itertools
 import unittest
+import pytest
 from unittest.mock import patch, MagicMock, Mock
 from ssl import CERT_NONE, CERT_REQUIRED
 
-import pyarrow
-
 import databricks.sql
 from databricks.sql import utils
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql import *
 from databricks.sql.auth.authenticators import AuthProvider
 from databricks.sql.thrift_backend import ThriftBackend
+from tests.e2e.predicate import pysql_supports_arrow
 
+try:
+    import pyarrow
+except ImportError:
+    pyarrow = None
 
 def retry_policy_factory():
     return {  # (type, default, min, max)
@@ -24,7 +28,7 @@ def retry_policy_factory():
         "_retry_delay_default": (float, 5, 1, 60),
     }
 
-
+@pytest.mark.skipif(not pysql_supports_arrow(), reason="Skipping because pyarrow is not installed")
 class ThriftBackendTestSuite(unittest.TestCase):
     okay_status = ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS)