From 8b167b9d7fc030eb7eae7b4ae59416f44b7ae7f6 Mon Sep 17 00:00:00 2001
From: Jothi Prakash <jothi.prakash@databricks.com>
Date: Wed, 18 Sep 2024 22:40:17 +0530
Subject: [PATCH] Refractor

---
 .../tests/unit/test_fetches_bench.py                   |  7 ++++++-
 .../tests/unit/test_thrift_backend.py                  | 10 +++++++---
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/databricks_sql_connector_core/tests/unit/test_fetches_bench.py b/databricks_sql_connector_core/tests/unit/test_fetches_bench.py
index e322b44a..7b07d473 100644
--- a/databricks_sql_connector_core/tests/unit/test_fetches_bench.py
+++ b/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.
diff --git a/databricks_sql_connector_core/tests/unit/test_thrift_backend.py b/databricks_sql_connector_core/tests/unit/test_thrift_backend.py
index 4bcf84d2..bd7ebac0 100644
--- a/databricks_sql_connector_core/tests/unit/test_thrift_backend.py
+++ b/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)