add more sea e2e testing

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

Author: varun-edachali-dbx

tests/e2e/common/decimal_tests.py

@@ -38,22 +38,42 @@ class DecimalTestsMixin:
         ),
     ]
 
+    @pytest.mark.parametrize(
+        "backend_params",
+        [
+            {},
+            {
+                "use_sea": True,
+            },
+        ],
+    )
     @pytest.mark.parametrize(
         "decimal, expected_value, expected_type", decimal_and_expected_results
     )
-    def test_decimals(self, decimal, expected_value, expected_type):
-        with self.cursor({}) as cursor:
+    def test_decimals(self, decimal, expected_value, expected_type, backend_params):
+        with self.cursor(backend_params) as cursor:
             query = "SELECT CAST ({})".format(decimal)
             cursor.execute(query)
             table = cursor.fetchmany_arrow(1)
             assert table.field(0).type == expected_type
             assert table.to_pydict().popitem()[1][0] == expected_value
 
+    @pytest.mark.parametrize(
+        "backend_params",
+        [
+            {},
+            {
+                "use_sea": True,
+            },
+        ],
+    )
     @pytest.mark.parametrize(
         "decimals, expected_values, expected_type", multi_decimals_and_expected_results
     )
-    def test_multi_decimals(self, decimals, expected_values, expected_type):
-        with self.cursor({}) as cursor:
+    def test_multi_decimals(
+        self, decimals, expected_values, expected_type, backend_params
+    ):
+        with self.cursor(backend_params) as cursor:
             union_str = " UNION ".join(
                 ["(SELECT CAST ({}))".format(dec) for dec in decimals]
             )

tests/e2e/common/large_queries_mixin.py

@@ -42,7 +42,16 @@ def fetch_rows(self, cursor, row_count, fetchmany_size):
             + "assuming 10K fetch size."
         )
 
-    def test_query_with_large_wide_result_set(self):
+    @pytest.mark.parametrize(
+        "backend_params",
+        [
+            {},
+            {
+                "use_sea": True,
+            },
+        ],
+    )
+    def test_query_with_large_wide_result_set(self, backend_params):
         resultSize = 300 * 1000 * 1000  # 300 MB
         width = 8192  # B
         rows = resultSize // width
@@ -52,7 +61,7 @@ def test_query_with_large_wide_result_set(self):
         fetchmany_size = 10 * 1024 * 1024 // width
         # This is used by PyHive tests to determine the buffer size
         self.arraysize = 1000
-        with self.cursor() as cursor:
+        with self.cursor(backend_params) as cursor:
             for lz4_compression in [False, True]:
                 cursor.connection.lz4_compression = lz4_compression
                 uuids = ", ".join(["uuid() uuid{}".format(i) for i in range(cols)])
@@ -68,7 +77,16 @@ def test_query_with_large_wide_result_set(self):
                     assert row[0] == row_id  # Verify no rows are dropped in the middle.
                     assert len(row[1]) == 36
 
-    def test_query_with_large_narrow_result_set(self):
+    @pytest.mark.parametrize(
+        "backend_params",
+        [
+            {},
+            {
+                "use_sea": True,
+            },
+        ],
+    )
+    def test_query_with_large_narrow_result_set(self, backend_params):
         resultSize = 300 * 1000 * 1000  # 300 MB
         width = 8  # sizeof(long)
         rows = resultSize / width
@@ -77,12 +95,21 @@ def test_query_with_large_narrow_result_set(self):
         fetchmany_size = 10 * 1024 * 1024 // width
         # This is used by PyHive tests to determine the buffer size
         self.arraysize = 10000000
-        with self.cursor() as cursor:
+        with self.cursor(backend_params) as cursor:
             cursor.execute("SELECT * FROM RANGE({rows})".format(rows=rows))
             for row_id, row in enumerate(self.fetch_rows(cursor, rows, fetchmany_size)):
                 assert row[0] == row_id
 
-    def test_long_running_query(self):
+    @pytest.mark.parametrize(
+        "backend_params",
+        [
+            {},
+            {
+                "use_sea": True,
+            },
+        ],
+    )
+    def test_long_running_query(self, backend_params):
         """Incrementally increase query size until it takes at least 3 minutes,
         and asserts that the query completes successfully.
         """
@@ -92,7 +119,7 @@ def test_long_running_query(self):
         duration = -1
         scale0 = 10000
         scale_factor = 1
-        with self.cursor() as cursor:
+        with self.cursor(backend_params) as cursor:
             while duration < min_duration:
                 assert scale_factor < 1024, "Detected infinite loop"
                 start = time.time()

tests/e2e/common/retry_test_mixins.py

@@ -17,17 +17,37 @@
 
 
 class Client429ResponseMixin:
-    def test_client_should_retry_automatically_when_getting_429(self):
-        with self.cursor() as cursor:
+    @pytest.mark.parametrize(
+        "backend_params",
+        [
+            {},
+            {
+                "use_sea": True,
+            },
+        ],
+    )
+    def test_client_should_retry_automatically_when_getting_429(self, backend_params):
+        with self.cursor(backend_params) as cursor:
             for _ in range(10):
                 cursor.execute("SELECT 1")
                 rows = cursor.fetchall()
                 self.assertEqual(len(rows), 1)
                 self.assertEqual(rows[0][0], 1)
 
-    def test_client_should_not_retry_429_if_RateLimitRetry_is_0(self):
+    @pytest.mark.parametrize(
+        "backend_params",
+        [
+            {},
+            {
+                "use_sea": True,
+            },
+        ],
+    )
+    def test_client_should_not_retry_429_if_RateLimitRetry_is_0(self, backend_params):
         with pytest.raises(self.error_type) as cm:
-            with self.cursor(self.conf_to_disable_rate_limit_retries) as cursor:
+            with self.cursor(
+                {**self.conf_to_disable_rate_limit_retries, **backend_params}
+            ) as cursor:
                 for _ in range(10):
                     cursor.execute("SELECT 1")
                     rows = cursor.fetchall()
@@ -46,14 +66,39 @@ def test_client_should_not_retry_429_if_RateLimitRetry_is_0(self):
 
 
 class Client503ResponseMixin:
-    def test_wait_cluster_startup(self):
-        with self.cursor() as cursor:
+    @pytest.mark.parametrize(
+        "backend_params",
+        [
+            {},
+            {
+                "use_sea": True,
+            },
+        ],
+    )
+    def test_wait_cluster_startup(self, backend_params):
+        with self.cursor(backend_params) as cursor:
             cursor.execute("SELECT 1")
             cursor.fetchall()
 
-    def _test_retry_disabled_with_message(self, error_msg_substring, exception_type):
+    @pytest.mark.parametrize(
+        "backend_params",
+        [
+            {},
+            {
+                "use_sea": True,
+            },
+        ],
+    )
+    def _test_retry_disabled_with_message(
+        self, error_msg_substring, exception_type, backend_params
+    ):
         with pytest.raises(exception_type) as cm:
-            with self.connection(self.conf_to_disable_temporarily_unavailable_retries):
+            with self.connection(
+                {
+                    **self.conf_to_disable_temporarily_unavailable_retries,
+                    **backend_params,
+                }
+            ):
                 pass
         assert error_msg_substring in str(cm.exception)