Test against Spark Connect server, add to CI

G-Research · Aug 1, 2024 · 152bbff · 152bbff
1 parent 4ba4159
commit 152bbff
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 3 deletions.
diff --git a/.github/actions/test-python/action.yml b/.github/actions/test-python/action.yml
@@ -112,7 +112,22 @@ runs:
       find python/test -name 'test*.py' > tests
       while read test
       do
-        if ! $SPARK_HOME/bin/spark-submit --master "local[2]" --packages uk.co.gresearch.spark:spark-extension_${{ inputs.scala-compat-version }}:$SPARK_EXTENSION_VERSION "$test" test-results
+        if ! $SPARK_HOME/bin/spark-submit --master "local[2]" --packages uk.co.gresearch.spark:spark-extension_${{ inputs.scala-compat-version }}:$SPARK_EXTENSION_VERSION "$test" test-results-submit
+        then
+          state="fail"
+        fi
+      done < tests
+      if [[ "$state" == "fail" ]]; then exit 1; fi
+    shell: bash
+
+  - name: Python Integration Tests (Spark Connect)
+    env:
+      PYTHONPATH: python:python/test
+    run: |
+      find python/test -name 'test*.py' > tests
+      while read test
+      do
+        if ! $SPARK_HOME/bin/spark-submit --master "local[2]" --packages uk.co.gresearch.spark:spark-extension_${{ inputs.scala-compat-version }}:$SPARK_EXTENSION_VERSION "$test" test-results-connect
         then
           state="fail"
         fi
@@ -135,7 +150,10 @@ runs:
     uses: actions/upload-artifact@v4
     with:
       name: Python Test Results (Spark ${{ inputs.spark-version }} Scala ${{ inputs.scala-version }} Python ${{ inputs.python-version }})
-      path: test-results/*.xml
+      path: |
+        test-results/*.xml
+        test-results-submit/*.xml
+        test-results-connect/*.xml
 
 branding:
   icon: 'check-circle'

diff --git a/python/test/spark_common.py b/python/test/spark_common.py
@@ -80,7 +80,9 @@ def get_spark_config(path) -> SparkConf:
     def get_spark_session(cls) -> SparkSession:
         builder = SparkSession.builder
 
-        if 'PYSPARK_GATEWAY_PORT' in os.environ:
+        if 'TEST_SPARK_CONNECT_SERVER' in os.environ:
+            builder.remote(os.environ['TEST_SPARK_CONNECT_SERVER'])
+        elif 'PYSPARK_GATEWAY_PORT' in os.environ:
             logging.info('Running inside existing Spark environment')
         else:
             logging.info('Setting up Spark environment')