test: Reduce execution time for the unit tests (#2720)

aws · Mar 12, 2024 · e8a5a96 · e8a5a96
1 parent e2c960b
commit e8a5a96
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/tests/unit/test_athena.py b/tests/unit/test_athena.py
@@ -1426,10 +1426,10 @@ def test_athena_generate_create_query(path, glue_database, glue_table):
     assert query == create_view
 
 
-def test_get_query_execution(workgroup0, workgroup1):
-    query_execution_ids = wr.athena.list_query_executions(workgroup=workgroup0) + wr.athena.list_query_executions(
-        workgroup=workgroup1
-    )
+def test_get_query_execution(workgroup0: str, workgroup1: str):
+    query_execution_ids = wr.athena.list_query_executions(
+        workgroup=workgroup0, max_results=10
+    ) + wr.athena.list_query_executions(workgroup=workgroup1, max_results=10)
     assert query_execution_ids
     query_execution_detail = wr.athena.get_query_execution(query_execution_id=query_execution_ids[0])
     query_executions_df = wr.athena.get_query_executions(query_execution_ids)
@@ -1450,7 +1450,7 @@ def test_list_query_executions_max_results(workgroup0: str, max_results: int):
     for _ in range(max_results + 1):
         wr.athena.start_query_execution(sql="SELECT random(10)", workgroup=workgroup0, wait=False)
 
-    query_execution_ids = wr.athena.list_query_executions(workgroup=workgroup0)
+    query_execution_ids = wr.athena.list_query_executions(workgroup=workgroup0, max_results=max_results + 1)
     assert len(query_execution_ids) > max_results
 
     query_execution_ids_max_results = wr.athena.list_query_executions(workgroup=workgroup0, max_results=max_results)

diff --git a/tox.ini b/tox.ini
@@ -18,6 +18,7 @@ commands =
        pytest -n {posargs} -s -v --timeout=300 --reruns=2 --reruns-delay=15 \
               --cov=awswrangler --cov-report=xml --cov-report term-missing --cov-branch \
               --cov-fail-under={env:COV_FAIL_UNDER} \
+              --dist load --maxschedchunk 2 \
               --junitxml=test-reports/junit.xml --log-file=test-reports/logs.txt tests/unit
 
 [testenv:py{37,38,39,310,311}-distributed]