SEA: e2e testing

src/databricks/sql/types.py

@@ -187,6 +187,7 @@ def __contains__(self, item: Any) -> bool:
     # let object acts like class
     def __call__(self, *args: Any) -> "Row":
         """create new Row object"""
+
         if len(args) > len(self):
             raise ValueError(
                 "Can not create Row with fields %s, expected %d values "
@@ -229,13 +230,15 @@ def __reduce__(
         self,
     ) -> Union[str, Tuple[Any, ...]]:
         """Returns a tuple so Python knows how to pickle Row."""
+
         if hasattr(self, "__fields__"):
             return (_create_row, (self.__fields__, tuple(self)))
         else:
             return tuple.__reduce__(self)
 
     def __repr__(self) -> str:
         """Printable representation of Row used in Python REPL."""
+
         if hasattr(self, "__fields__"):
             return "Row(%s)" % ", ".join(
                 "%s=%r" % (k, v) for k, v in zip(self.__fields__, tuple(self))

src/databricks/sql/utils.py

@@ -740,7 +740,6 @@ def convert_decimals_in_arrow_table(table, description) -> "pyarrow.Table":
 
 
 def convert_to_assigned_datatypes_in_column_table(column_table, description):
-
     converted_column_table = []
     for i, col in enumerate(column_table):
         if description[i][1] == "decimal":

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/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)