Skip to content

Commit c621c0c

Browse files
introduce SEA cloudfetch e2e tests
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent f84578a commit c621c0c

File tree

1 file changed

+49
-4
lines changed

1 file changed

+49
-4
lines changed

tests/e2e/test_driver.py

Lines changed: 49 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -180,10 +180,19 @@ def test_cloud_fetch(self):
180180

181181

182182
class TestPySQLAsyncQueriesSuite(PySQLPytestTestCase):
183-
def test_execute_async__long_running(self):
183+
@pytest.mark.parametrize(
184+
"extra_params",
185+
[
186+
{},
187+
{
188+
"use_sea": True,
189+
},
190+
],
191+
)
192+
def test_execute_async__long_running(self, extra_params):
184193

185194
long_running_query = "SELECT COUNT(*) FROM RANGE(10000 * 16) x JOIN RANGE(10000) y ON FROM_UNIXTIME(x.id * y.id, 'yyyy-MM-dd') LIKE '%not%a%date%'"
186-
with self.cursor() as cursor:
195+
with self.cursor(extra_params) as cursor:
187196
cursor.execute_async(long_running_query)
188197

189198
## Polling after every POLLING_INTERVAL seconds
@@ -226,7 +235,16 @@ def test_execute_async__small_result(self, extra_params):
226235

227236
assert result[0].asDict() == {"1": 1}
228237

229-
def test_execute_async__large_result(self):
238+
@pytest.mark.parametrize(
239+
"extra_params",
240+
[
241+
{},
242+
{
243+
"use_sea": True,
244+
},
245+
],
246+
)
247+
def test_execute_async__large_result(self, extra_params):
230248
x_dimension = 1000
231249
y_dimension = 1000
232250
large_result_query = f"""
@@ -240,7 +258,7 @@ def test_execute_async__large_result(self):
240258
RANGE({y_dimension}) y
241259
"""
242260

243-
with self.cursor() as cursor:
261+
with self.cursor(extra_params) as cursor:
244262
cursor.execute_async(large_result_query)
245263

246264
## Fake sleep for 5 secs
@@ -348,6 +366,9 @@ def test_incorrect_query_throws_exception(self):
348366
"use_cloud_fetch": False,
349367
"enable_query_result_lz4_compression": False,
350368
},
369+
{
370+
"use_sea": True,
371+
}
351372
],
352373
)
353374
def test_create_table_will_return_empty_result_set(self, extra_params):
@@ -558,6 +579,9 @@ def test_get_catalogs(self):
558579
"use_cloud_fetch": False,
559580
"enable_query_result_lz4_compression": False,
560581
},
582+
{
583+
"use_sea": True,
584+
}
561585
],
562586
)
563587
def test_get_arrow(self, extra_params):
@@ -631,6 +655,9 @@ def execute_really_long_query():
631655
"use_cloud_fetch": False,
632656
"enable_query_result_lz4_compression": False,
633657
},
658+
{
659+
"use_sea": True,
660+
}
634661
],
635662
)
636663
def test_can_execute_command_after_failure(self, extra_params):
@@ -653,6 +680,9 @@ def test_can_execute_command_after_failure(self, extra_params):
653680
"use_cloud_fetch": False,
654681
"enable_query_result_lz4_compression": False,
655682
},
683+
{
684+
"use_sea": True,
685+
}
656686
],
657687
)
658688
def test_can_execute_command_after_success(self, extra_params):
@@ -677,6 +707,9 @@ def generate_multi_row_query(self):
677707
"use_cloud_fetch": False,
678708
"enable_query_result_lz4_compression": False,
679709
},
710+
{
711+
"use_sea": True,
712+
}
680713
],
681714
)
682715
def test_fetchone(self, extra_params):
@@ -721,6 +754,9 @@ def test_fetchall(self, extra_params):
721754
"use_cloud_fetch": False,
722755
"enable_query_result_lz4_compression": False,
723756
},
757+
{
758+
"use_sea": True,
759+
}
724760
],
725761
)
726762
def test_fetchmany_when_stride_fits(self, extra_params):
@@ -741,6 +777,9 @@ def test_fetchmany_when_stride_fits(self, extra_params):
741777
"use_cloud_fetch": False,
742778
"enable_query_result_lz4_compression": False,
743779
},
780+
{
781+
"use_sea": True,
782+
}
744783
],
745784
)
746785
def test_fetchmany_in_excess(self, extra_params):
@@ -761,6 +800,9 @@ def test_fetchmany_in_excess(self, extra_params):
761800
"use_cloud_fetch": False,
762801
"enable_query_result_lz4_compression": False,
763802
},
803+
{
804+
"use_sea": True,
805+
}
764806
],
765807
)
766808
def test_iterator_api(self, extra_params):
@@ -846,6 +888,9 @@ def test_timestamps_arrow(self):
846888
"use_cloud_fetch": False,
847889
"enable_query_result_lz4_compression": False,
848890
},
891+
{
892+
"use_sea": True,
893+
}
849894
],
850895
)
851896
def test_multi_timestamps_arrow(self, extra_params):

0 commit comments

Comments
 (0)