From d509a60cc0c8ee3f7cc72b1ef1c3d30e6cbb785e Mon Sep 17 00:00:00 2001
From: Gera Shegalov <gera@apache.org>
Date: Fri, 26 Mar 2021 07:12:48 -0700
Subject: [PATCH] Refactor logic for parallel testing (#2016)

Signed-off-by: Gera Shegalov <gera@apache.org>
---
 integration_tests/run_pyspark_from_build.sh | 60 ++++++++++-----------
 1 file changed, 27 insertions(+), 33 deletions(-)

diff --git a/integration_tests/run_pyspark_from_build.sh b/integration_tests/run_pyspark_from_build.sh
index ad7d8c7b180..8d920e6eb48 100755
--- a/integration_tests/run_pyspark_from_build.sh
+++ b/integration_tests/run_pyspark_from_build.sh
@@ -84,11 +84,11 @@ else
     then
         # With xdist 0 and 1 are the same parallelsm but
         # 0 is more effecient
-        TEST_PARALLEL_OPTS=""
+        TEST_PARALLEL_OPTS=()
         MEMORY_FRACTION='1'
     else
         MEMORY_FRACTION=`python -c "print(1/($TEST_PARALLEL + 1))"`
-        TEST_PARALLEL_OPTS="-n $TEST_PARALLEL"
+        TEST_PARALLEL_OPTS=("-n" "$TEST_PARALLEL")
     fi
     RUN_DIR="$SCRIPTPATH"/target/run_dir
     mkdir -p "$RUN_DIR"
@@ -99,41 +99,35 @@ else
     ## Under cloud environment, overwrite the '--std_input_path' param to point to the distributed file path
     INPUT_PATH=${INPUT_PATH:-"$SCRIPTPATH"}
 
-    if [[ "${TEST_PARALLEL_OPTS}" != "" ]];
+    RUN_TESTS_COMMAND=("$SCRIPTPATH"/runtests.py
+      --rootdir
+      "$LOCAL_ROOTDIR"
+      "$LOCAL_ROOTDIR"/src/main/python)
+
+    TEST_COMMON_OPTS=(-v
+          -rfExXs
+          "$TEST_TAGS"
+          --std_input_path="$INPUT_PATH"/src/test/resources
+          --color=yes
+          $TEST_TYPE_PARAM
+          "$TEST_ARGS"
+          $RUN_TEST_PARAMS
+          "$@")
+
+    export PYSP_TEST_spark_driver_extraClassPath="${ALL_JARS// /:}"
+    export PYSP_TEST_spark_driver_extraJavaOptions="-ea -Duser.timezone=UTC $COVERAGE_SUBMIT_FLAGS"
+    export PYSP_TEST_spark_executor_extraJavaOptions='-ea -Duser.timezone=UTC'
+    export PYSP_TEST_spark_ui_showConsoleProgress='false'
+    export PYSP_TEST_spark_sql_session_timeZone='UTC'
+    export PYSP_TEST_spark_sql_shuffle_partitions='12'
+    if ((${#TEST_PARALLEL_OPTS[@]} > 0));
     then
-        export PYSP_TEST_spark_driver_extraClassPath="${ALL_JARS// /:}"
-        export PYSP_TEST_spark_driver_extraJavaOptions="-ea -Duser.timezone=UTC $COVERAGE_SUBMIT_FLAGS"
-        export PYSP_TEST_spark_executor_extraJavaOptions='-ea -Duser.timezone=UTC'
-        export PYSP_TEST_spark_ui_showConsoleProgress='false'
-        export PYSP_TEST_spark_sql_session_timeZone='UTC'
-        export PYSP_TEST_spark_sql_shuffle_partitions='12'
         export PYSP_TEST_spark_rapids_memory_gpu_allocFraction=$MEMORY_FRACTION
         export PYSP_TEST_spark_rapids_memory_gpu_maxAllocFraction=$MEMORY_FRACTION
-
-        python \
-          "$SCRIPTPATH"/runtests.py --rootdir "$LOCAL_ROOTDIR" "$LOCAL_ROOTDIR"/src/main/python \
-          $TEST_PARALLEL_OPTS \
-          -v -rfExXs "$TEST_TAGS" \
-          --std_input_path="$INPUT_PATH"/src/test/resources/ \
-          --color=yes \
-          $TEST_TYPE_PARAM \
-          "$TEST_ARGS" \
-          $RUN_TEST_PARAMS \
-          "$@"
+        python "${RUN_TESTS_COMMAND[@]}" "${TEST_PARALLEL_OPTS[@]}" "${TEST_COMMON_OPTS[@]}"
     else
         "$SPARK_HOME"/bin/spark-submit --jars "${ALL_JARS// /,}" \
-          --conf "spark.driver.extraJavaOptions=-ea -Duser.timezone=UTC $COVERAGE_SUBMIT_FLAGS" \
-          --conf 'spark.executor.extraJavaOptions=-ea -Duser.timezone=UTC' \
-          --conf 'spark.sql.session.timeZone=UTC' \
-          --conf 'spark.sql.shuffle.partitions=12' \
-          $SPARK_SUBMIT_FLAGS \
-          "$SCRIPTPATH"/runtests.py --rootdir "$LOCAL_ROOTDIR" "$LOCAL_ROOTDIR"/src/main/python \
-          -v -rfExXs "$TEST_TAGS" \
-          --std_input_path="$INPUT_PATH"/src/test/resources/ \
-          --color=yes \
-          $TEST_TYPE_PARAM \
-          "$TEST_ARGS" \
-          $RUN_TEST_PARAMS \
-          "$@"
+            --driver-java-options "$PYSP_TEST_spark_driver_extraJavaOptions" \
+            $SPARK_SUBMIT_FLAGS "${RUN_TESTS_COMMAND[@]}" "${TEST_COMMON_OPTS[@]}"
     fi
 fi